digitalkin 0.3.1.dev2__py3-none-any.whl → 0.3.2a3__py3-none-any.whl

monitoring/tests/test_metrics.py

@@ -0,0 +1,172 @@
+"""Tests for metrics collection.
+
+Run with: python -m pytest tests/test_metrics.py
+"""
+
+import sys
+from pathlib import Path
+
+import pytest
+
+# Add the parent directory to the path so we can import digitalkin_observability
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from digitalkin_observability import MetricsCollector, PrometheusExporter, get_metrics
+
+
+class TestMetricsCollector:
+    """Tests for MetricsCollector singleton."""
+
+    def setup_method(self) -> None:
+        """Reset metrics before each test."""
+        get_metrics().reset()
+
+    def test_singleton_returns_same_instance(self) -> None:
+        """Test that get_metrics returns the same instance."""
+        m1 = get_metrics()
+        m2 = get_metrics()
+        assert m1 is m2
+
+    def test_inc_jobs_started(self) -> None:
+        """Test incrementing jobs started counter."""
+        metrics = get_metrics()
+        metrics.inc_jobs_started("TestModule")
+
+        assert metrics.jobs_started_total == 1
+        assert metrics.active_jobs == 1
+
+    def test_inc_jobs_completed(self) -> None:
+        """Test incrementing jobs completed counter."""
+        metrics = get_metrics()
+        metrics.inc_jobs_started("TestModule")
+        metrics.inc_jobs_completed("TestModule", 1.5)
+
+        assert metrics.jobs_completed_total == 1
+        assert metrics.active_jobs == 0
+        assert metrics.job_duration_seconds.count == 1
+        assert metrics.job_duration_seconds.total_sum == 1.5
+
+    def test_inc_jobs_failed(self) -> None:
+        """Test incrementing jobs failed counter."""
+        metrics = get_metrics()
+        metrics.inc_jobs_started("TestModule")
+        metrics.inc_jobs_failed("TestModule")
+
+        assert metrics.jobs_failed_total == 1
+        assert metrics.active_jobs == 0
+
+    def test_inc_jobs_cancelled(self) -> None:
+        """Test incrementing jobs cancelled counter."""
+        metrics = get_metrics()
+        metrics.inc_jobs_started("TestModule")
+        metrics.inc_jobs_cancelled("TestModule")
+
+        assert metrics.jobs_cancelled_total == 1
+        assert metrics.active_jobs == 0
+
+    def test_inc_messages_sent(self) -> None:
+        """Test incrementing messages sent counter."""
+        metrics = get_metrics()
+        metrics.inc_messages_sent("message")
+        metrics.inc_messages_sent("file")
+        metrics.inc_messages_sent()
+
+        assert metrics.messages_sent_total == 3
+
+    def test_queue_depth_tracking(self) -> None:
+        """Test queue depth tracking."""
+        metrics = get_metrics()
+        metrics.set_queue_depth("job1", 5)
+        metrics.set_queue_depth("job2", 3)
+
+        assert metrics.queue_depth["job1"] == 5
+        assert metrics.queue_depth["job2"] == 3
+
+        metrics.clear_queue_depth("job1")
+        assert "job1" not in metrics.queue_depth
+
+    def test_snapshot(self) -> None:
+        """Test snapshot returns all metrics."""
+        metrics = get_metrics()
+        metrics.inc_jobs_started("TestModule")
+        metrics.inc_jobs_completed("TestModule", 0.5)
+        metrics.inc_messages_sent("message")
+
+        snapshot = metrics.snapshot()
+
+        assert snapshot["jobs_started_total"] == 1
+        assert snapshot["jobs_completed_total"] == 1
+        assert snapshot["messages_sent_total"] == 1
+        assert "job_duration_seconds" in snapshot
+        assert "by_module" in snapshot
+        assert "TestModule" in snapshot["by_module"]
+
+    def test_histogram_observe(self) -> None:
+        """Test histogram observations."""
+        metrics = get_metrics()
+        metrics.observe_grpc_duration(0.05)
+        metrics.observe_grpc_duration(0.15)
+
+        assert metrics.grpc_request_duration_seconds.count == 2
+        assert metrics.grpc_request_duration_seconds.total_sum == pytest.approx(0.2)
+
+    def test_reset_clears_all_metrics(self) -> None:
+        """Test reset clears all metrics."""
+        metrics = get_metrics()
+        metrics.inc_jobs_started("TestModule")
+        metrics.inc_errors()
+
+        metrics.reset()
+
+        assert metrics.jobs_started_total == 0
+        assert metrics.errors_total == 0
+        assert metrics.active_jobs == 0
+
+
+class TestPrometheusExporter:
+    """Tests for Prometheus exporter."""
+
+    def setup_method(self) -> None:
+        """Reset metrics before each test."""
+        get_metrics().reset()
+
+    def test_export_returns_string(self) -> None:
+        """Test that export returns a string."""
+        output = PrometheusExporter.export()
+        assert isinstance(output, str)
+
+    def test_export_contains_job_counters(self) -> None:
+        """Test export contains job counters."""
+        metrics = get_metrics()
+        metrics.inc_jobs_started("TestModule")
+
+        output = PrometheusExporter.export()
+
+        assert "digitalkin_jobs_started_total 1" in output
+        assert "digitalkin_active_jobs 1" in output
+
+    def test_export_contains_histogram(self) -> None:
+        """Test export contains histogram data."""
+        metrics = get_metrics()
+        metrics.observe_grpc_duration(0.05)
+
+        output = PrometheusExporter.export()
+
+        assert "digitalkin_grpc_request_duration_seconds" in output
+        assert "# TYPE digitalkin_grpc_request_duration_seconds histogram" in output
+
+    def test_export_contains_module_breakdown(self) -> None:
+        """Test export contains per-module breakdown."""
+        metrics = get_metrics()
+        metrics.inc_jobs_started("MyModule")
+
+        output = PrometheusExporter.export()
+
+        assert 'digitalkin_jobs_by_module{module="MyModule",status="started"} 1' in output
+
+    def test_export_contains_help_and_type(self) -> None:
+        """Test export contains HELP and TYPE comments."""
+        output = PrometheusExporter.export()
+
+        assert "# HELP digitalkin_jobs_started_total" in output
+        assert "# TYPE digitalkin_jobs_started_total counter" in output

services/filesystem_module.py

@@ -3,16 +3,16 @@
 import asyncio
 import datetime
 from collections.abc import Callable
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 from pydantic import BaseModel, Field
 
 from digitalkin.logger import logger
 from digitalkin.models.module import ModuleStatus
 from digitalkin.modules.archetype_module import ArchetypeModule
+from digitalkin.services.filesystem.filesystem_strategy import FileFilter, UploadFileData
 from digitalkin.services.services_config import ServicesConfig
 from digitalkin.services.services_models import ServicesMode
-from digitalkin.services.filesystem.filesystem_strategy import FilesystemRecord, FileFilter, UploadFileData
 
 
 class ExampleInput(BaseModel):
@@ -62,7 +62,10 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
 
     # Define services_config_params with default values
     services_config_strategies = {}
-    services_config_params = {"cost": {"config": {}}, "storage": {"config": {}}} # Filesystem has no config but it's enabled
+    services_config_params = {
+        "cost": {"config": {}},
+        "storage": {"config": {}},
+    }  # Filesystem has no config but it's enabled
 
     def __init__(self, job_id: str, mission_id: str, setup_version_id: str) -> None:
         """Initialize the example module.
@@ -128,7 +131,6 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
             callback(record.model_dump())
         # Call the callback with the output data
 
-
         # Wait a bit to simulate processing time
         await asyncio.sleep(1)
 
@@ -173,7 +175,7 @@ async def test_module() -> None:
 
     # Check the storage
     if module.status == ModuleStatus.STOPPED:
-        files, nb_results = module.filesystem.get_files(
+        files, _nb_results = module.filesystem.get_files(
             filters=FileFilter(name="example_output.txt", context="test-mission-123"),
         )
         for file in files:

services/storage_module.py

@@ -64,7 +64,7 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
 
     # Define services_config_params with default values
     services_config_strategies = {}
-    services_config_params = {"storage": {"config": {"example": ExampleOutput}},"cost": {"config":{}}}
+    services_config_params = {"storage": {"config": {"example": ExampleOutput}}, "cost": {"config": {}}}
 
     def __init__(self, job_id: str, mission_id: str, setup_version_id: str) -> None:
         """Initialize the example module.
@@ -184,7 +184,9 @@ async def test_module() -> None:
 def test_storage_directly() -> None:
     """Test the storage service directly."""
     # Initialize storage service
-    storage = ServicesConfig().storage(mission_id="test-mission",setup_version_id="test-setup-123", config={"example": ExampleStorage})
+    storage = ServicesConfig().storage(
+        mission_id="test-mission", setup_version_id="test-setup-123", config={"example": ExampleStorage}
+    )
 
     # Create a test record
     storage.store("example", "test_table", {"test_key": "test_value"}, "OUTPUT")

digitalkin/grpc_servers/registry_server.py

@@ -1,65 +0,0 @@
-"""Registry gRPC server implementation for DigitalKin."""
-
-from digitalkin_proto.agentic_mesh_protocol.module_registry.v1 import (
-    module_registry_service_pb2,
-    module_registry_service_pb2_grpc,
-)
-
-from digitalkin.grpc_servers._base_server import BaseServer
-from digitalkin.grpc_servers.registry_servicer import RegistryModule, RegistryServicer
-from digitalkin.logger import logger
-from digitalkin.models.grpc_servers.models import RegistryServerConfig
-
-
-class RegistryServer(BaseServer):
-    """gRPC server for DigitalKin module registry.
-
-    This server implements the ModuleRegistryService which allows modules to register
-    themselves and be discovered by other components in the system.
-
-    Attributes:
-        config: Server configuration.
-        registry_servicer: The gRPC servicer handling registry requests.
-    """
-
-    def __init__(
-        self,
-        config: RegistryServerConfig,
-    ) -> None:
-        """Initialize the registry server.
-
-        Args:
-            config: Server configuration.
-        """
-        super().__init__(config)
-        self.config = config
-        self.registry_servicer: RegistryServicer | None = None
-
-    def _register_servicers(self) -> None:
-        """Register the registry servicer with the gRPC server.
-
-        Raises:
-            RuntimeError: If server is not registered during server creation
-        """
-        if self.server is None:
-            msg = "Server must be created before registering servicers"
-            raise RuntimeError(msg)
-
-        logger.debug("Registering registry servicer")
-        self.registry_servicer = RegistryServicer()
-        self.register_servicer(
-            self.registry_servicer,
-            module_registry_service_pb2_grpc.add_ModuleRegistryServiceServicer_to_server,
-            service_descriptor=module_registry_service_pb2.DESCRIPTOR,
-        )
-        logger.debug("Registered registry servicer")
-
-    def get_registered_modules(self) -> list[RegistryModule]:
-        """Get a list of all registered modules.
-
-        Returns:
-            A list of module information objects.
-        """
-        if self.registry_servicer:
-            return list(self.registry_servicer.registered_modules.values())
-        return []