RSDK-12467: Add and test historical module data (#1031)

Author: lia-viam

src/viam/module/resource_data_consumer.py

@@ -0,0 +1,35 @@
+import datetime
+import os
+from typing import Any, Dict, List
+
+from viam.app.viam_client import ViamClient
+
+
+class ResourceDataConsumer:
+    """Client for retrieving historical module data from app.
+
+    Inherit from this class in a module to get access to historical module data.
+    """
+
+    @classmethod
+    def construct_query(cls, part_id: str, resource_name: str, time_back: datetime.timedelta):
+        return [
+            {
+                "$match": {
+                    "part_id": part_id,
+                    "component_name": resource_name,
+                    "time_received": {"$gte": datetime.datetime.now() - time_back},
+                }
+            }
+        ]
+
+    @classmethod
+    async def query_tabular_data(cls, resource_name: str, time_back: datetime.timedelta, **kwargs) -> List[Dict[str, Any]]:
+        """Return historical data for this module, queried with MQL."""
+        viam_client = await ViamClient.create_from_env_vars()
+
+        org_id = os.environ["VIAM_PRIMARY_ORG_ID"]
+        part_id = os.environ["VIAM_MACHINE_PART_ID"]
+
+        query = cls.construct_query(part_id=part_id, resource_name=resource_name, time_back=time_back)
+        return await viam_client.data_client.tabular_data_by_mql(org_id, query)

tests/test_module.py

@@ -1,10 +1,14 @@
+import datetime
+import os
+import uuid
 from unittest import mock
 
 import pytest
 from grpclib.testing import ChannelFor
 
 from viam.errors import GRPCError
 from viam.module import Module
+from viam.module.resource_data_consumer import ResourceDataConsumer
 from viam.module.service import ModuleRPCService
 from viam.proto.app.robot import ComponentConfig
 from viam.proto.module import (
@@ -44,6 +48,33 @@ def service(module: Module) -> ModuleRPCService:
     return ModuleRPCService(module)
 
 
+class TestResourceDataConsumer:
+    async def test_historical_data(self):
+        with mock.patch("viam.app.data_client.DataClient.tabular_data_by_mql", new=mock.AsyncMock()) as mocked:
+            with mock.patch("viam.app.viam_client._get_access_token") as patched_auth:
+                ACCESS_TOKEN = "MY_ACCESS_TOKEN"
+                patched_auth.return_value = ACCESS_TOKEN
+
+                os.environ["VIAM_API_KEY"] = "MY_API_KEY"
+                os.environ["VIAM_API_KEY_ID"] = str(uuid.uuid4())
+                os.environ["VIAM_PRIMARY_ORG_ID"] = "my_org"
+                os.environ["VIAM_MACHINE_PART_ID"] = "my_part"
+
+                delta = datetime.timedelta(hours=2)
+
+                # Define a helper approx matcher because the time received fields will vary slightly
+                class DeltaApprox:
+                    def __eq__(self, other):
+                        gte = datetime.datetime.now() - delta
+                        return other - gte < datetime.timedelta(seconds=1)
+
+                query = ResourceDataConsumer.construct_query("my_part", "resource", delta)
+                query[0]["$match"]["time_received"]["$gte"] = DeltaApprox()
+
+                await ResourceDataConsumer.query_tabular_data("resource", delta)
+                mocked.assert_called_once_with("my_org", query)
+
+
 class TestModule:
     async def test_add_resource(self, module: Module):
         req = AddResourceRequest(