fix: add more test coverage and remove redundant command acknowledge messages in telemetry manager

ilana-n · ilana-n · commit 28ba6d4e7596 · 2025-10-19T18:22:35.000-07:00
diff --git a/tests/config/test_user_config.py b/tests/config/test_user_config.py
@@ -14,7 +14,7 @@
     TokenizerConfig,
     UserConfig,
 )
-from aiperf.common.enums import EndpointType
+from aiperf.common.enums import EndpointType, GPUTelemetryMode
 from aiperf.common.enums.timing_enums import TimingMode
 
 
@@ -201,3 +201,151 @@ def test_compute_artifact_directory(
 
     artifact_dir = config._compute_artifact_directory()
     assert artifact_dir == Path(expected_dir)
+
+
+@pytest.mark.parametrize(
+    "gpu_telemetry_input,expected_mode,expected_urls",
+    [
+        # No telemetry configured
+        ([], GPUTelemetryMode.SUMMARY, []),
+        # Dashboard mode only
+        (["dashboard"], GPUTelemetryMode.REALTIME_DASHBOARD, []),
+        # URLs only (no dashboard)
+        (
+            ["http://node1:9401/metrics"],
+            GPUTelemetryMode.SUMMARY,
+            ["http://node1:9401/metrics"],
+        ),
+        # Dashboard + URLs
+        (
+            ["dashboard", "http://node1:9401/metrics"],
+            GPUTelemetryMode.REALTIME_DASHBOARD,
+            ["http://node1:9401/metrics"],
+        ),
+        # Multiple URLs
+        (
+            ["http://node1:9401/metrics", "http://node2:9401/metrics"],
+            GPUTelemetryMode.SUMMARY,
+            ["http://node1:9401/metrics", "http://node2:9401/metrics"],
+        ),
+        # Dashboard + multiple URLs
+        (
+            [
+                "dashboard",
+                "http://node1:9401/metrics",
+                "http://node2:9401/metrics",
+            ],
+            GPUTelemetryMode.REALTIME_DASHBOARD,
+            ["http://node1:9401/metrics", "http://node2:9401/metrics"],
+        ),
+    ],
+)
+def test_parse_gpu_telemetry_config(gpu_telemetry_input, expected_mode, expected_urls):
+    """Test parsing of gpu_telemetry list into mode and URLs."""
+    config = UserConfig(
+        endpoint=EndpointConfig(
+            model_names=["test-model"],
+            type=EndpointType.CHAT,
+            custom_endpoint="test",
+        ),
+        gpu_telemetry=gpu_telemetry_input,
+    )
+
+    assert config.gpu_telemetry_mode == expected_mode
+    assert config.gpu_telemetry_urls == expected_urls
+
+
+def test_parse_gpu_telemetry_config_with_defaults():
+    """Test that gpu_telemetry_mode and gpu_telemetry_urls have correct defaults."""
+    config = UserConfig(
+        endpoint=EndpointConfig(
+            model_names=["test-model"],
+            type=EndpointType.CHAT,
+            custom_endpoint="test",
+        )
+    )
+
+    # Should have default values
+    assert config.gpu_telemetry_mode == GPUTelemetryMode.SUMMARY
+    assert config.gpu_telemetry_urls == []
+
+
+def test_parse_gpu_telemetry_config_preserves_existing_fields():
+    """Test that parsing GPU telemetry config doesn't affect other fields."""
+    config = UserConfig(
+        endpoint=EndpointConfig(
+            model_names=["test-model"],
+            type=EndpointType.CHAT,
+            custom_endpoint="test",
+            streaming=True,
+        ),
+        gpu_telemetry=["dashboard", "http://custom:9401/metrics"],
+    )
+
+    # Telemetry fields should be set
+    assert config.gpu_telemetry_mode == GPUTelemetryMode.REALTIME_DASHBOARD
+    assert config.gpu_telemetry_urls == ["http://custom:9401/metrics"]
+
+    # Other fields should be unchanged
+    assert config.endpoint.streaming is True
+    assert config.endpoint.model_names == ["test-model"]
+
+
+def test_gpu_telemetry_urls_extraction():
+    """Test that only http URLs are extracted from gpu_telemetry list."""
+    config = UserConfig(
+        endpoint=EndpointConfig(
+            model_names=["test-model"],
+            type=EndpointType.CHAT,
+            custom_endpoint="test",
+        ),
+        gpu_telemetry=[
+            "dashboard",  # Not a URL
+            "http://node1:9401/metrics",  # Valid URL
+            "https://node2:9401/metrics",  # Valid URL
+            "summary",  # Not a URL
+        ],
+    )
+
+    # Should extract only http/https URLs
+    assert len(config.gpu_telemetry_urls) == 2
+    assert "http://node1:9401/metrics" in config.gpu_telemetry_urls
+    assert "https://node2:9401/metrics" in config.gpu_telemetry_urls
+    assert "dashboard" not in config.gpu_telemetry_urls
+    assert "summary" not in config.gpu_telemetry_urls
+
+
+def test_gpu_telemetry_mode_detection():
+    """Test that dashboard mode is detected correctly in various positions."""
+    # Dashboard at beginning
+    config1 = UserConfig(
+        endpoint=EndpointConfig(
+            model_names=["test-model"],
+            type=EndpointType.CHAT,
+            custom_endpoint="test",
+        ),
+        gpu_telemetry=["dashboard", "http://node1:9401/metrics"],
+    )
+    assert config1.gpu_telemetry_mode == GPUTelemetryMode.REALTIME_DASHBOARD
+
+    # Dashboard at end
+    config2 = UserConfig(
+        endpoint=EndpointConfig(
+            model_names=["test-model"],
+            type=EndpointType.CHAT,
+            custom_endpoint="test",
+        ),
+        gpu_telemetry=["http://node1:9401/metrics", "dashboard"],
+    )
+    assert config2.gpu_telemetry_mode == GPUTelemetryMode.REALTIME_DASHBOARD
+
+    # No dashboard
+    config3 = UserConfig(
+        endpoint=EndpointConfig(
+            model_names=["test-model"],
+            type=EndpointType.CHAT,
+            custom_endpoint="test",
+        ),
+        gpu_telemetry=["http://node1:9401/metrics"],
+    )
+    assert config3.gpu_telemetry_mode == GPUTelemetryMode.SUMMARY
diff --git a/tests/mixins/test_realtime_telemetry_metrics_mixin.py b/tests/mixins/test_realtime_telemetry_metrics_mixin.py
@@ -0,0 +1,195 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from aiperf.common.config import ServiceConfig
+from aiperf.common.hooks import AIPerfHook
+from aiperf.common.messages import RealtimeTelemetryMetricsMessage
+from aiperf.common.mixins.realtime_telemetry_metrics_mixin import (
+    RealtimeTelemetryMetricsMixin,
+)
+from aiperf.common.models import MetricResult
+
+
+class TestRealtimeTelemetryMetricsMixin:
+    """Test suite for RealtimeTelemetryMetricsMixin functionality."""
+
+    @pytest.fixture
+    def mocked_mixin(self):
+        """Create a RealtimeTelemetryMetricsMixin instance with mocked dependencies."""
+        service_config = ServiceConfig()
+        mock_controller = MagicMock()
+
+        # Mock the MessageBusClientMixin.__init__ to avoid initialization issues
+        with patch(
+            "aiperf.common.mixins.message_bus_mixin.MessageBusClientMixin.__init__",
+            return_value=None,
+        ):
+            mixin = RealtimeTelemetryMetricsMixin(
+                service_config=service_config, controller=mock_controller
+            )
+            # Manually set attributes that would be set by parent __init__
+            mixin._controller = mock_controller
+            mixin._telemetry_metrics = []
+            mixin.run_hooks = AsyncMock()
+
+        return mixin
+
+    def test_mixin_initialization(self, mocked_mixin):
+        """Test that mixin initializes with correct attributes."""
+        assert hasattr(mocked_mixin, "_controller")
+        assert hasattr(mocked_mixin, "_telemetry_metrics")
+        assert hasattr(mocked_mixin, "_telemetry_metrics_lock")
+        assert mocked_mixin._telemetry_metrics == []
+
+    @pytest.mark.asyncio
+    async def test_on_realtime_telemetry_metrics_stores_metrics(self, mocked_mixin):
+        """Test that telemetry metrics are stored when message is received."""
+        metrics = [
+            MetricResult(tag="gpu_util", header="GPU Utilization", unit="%", avg=75.0),
+            MetricResult(
+                tag="gpu_memory", header="GPU Memory Used", unit="GB", avg=8.5
+            ),
+        ]
+
+        message = RealtimeTelemetryMetricsMessage(
+            service_id="records_manager", metrics=metrics
+        )
+
+        await mocked_mixin._on_realtime_telemetry_metrics(message)
+
+        # Verify metrics were stored
+        assert mocked_mixin._telemetry_metrics == metrics
+
+    @pytest.mark.asyncio
+    async def test_on_realtime_telemetry_metrics_triggers_hook(self, mocked_mixin):
+        """Test that receiving telemetry metrics triggers the appropriate hook."""
+
+        metrics = [
+            MetricResult(tag="gpu_util", header="GPU Utilization", unit="%", avg=75.0)
+        ]
+
+        message = RealtimeTelemetryMetricsMessage(
+            service_id="records_manager", metrics=metrics
+        )
+
+        await mocked_mixin._on_realtime_telemetry_metrics(message)
+
+        # Verify hook was triggered with correct arguments
+        mocked_mixin.run_hooks.assert_called_once_with(
+            AIPerfHook.ON_REALTIME_TELEMETRY_METRICS, metrics=metrics
+        )
+
+    @pytest.mark.asyncio
+    async def test_on_realtime_telemetry_metrics_replaces_previous_metrics(
+        self, mocked_mixin
+    ):
+        """Test that new metrics replace previous metrics (not append)."""
+        # Set initial metrics
+        initial_metrics = [
+            MetricResult(tag="old_metric", header="Old Metric", unit="ms", avg=10.0)
+        ]
+        mocked_mixin._telemetry_metrics = initial_metrics
+
+        # Receive new metrics
+        new_metrics = [
+            MetricResult(tag="new_metric", header="New Metric", unit="%", avg=50.0)
+        ]
+        message = RealtimeTelemetryMetricsMessage(
+            service_id="records_manager", metrics=new_metrics
+        )
+
+        await mocked_mixin._on_realtime_telemetry_metrics(message)
+
+        # Verify old metrics were replaced, not appended
+        assert mocked_mixin._telemetry_metrics == new_metrics
+        assert len(mocked_mixin._telemetry_metrics) == 1
+
+    @pytest.mark.asyncio
+    async def test_on_realtime_telemetry_metrics_with_empty_list(self, mocked_mixin):
+        """Test that receiving empty metrics list is handled correctly."""
+        message = RealtimeTelemetryMetricsMessage(
+            service_id="records_manager", metrics=[]
+        )
+
+        await mocked_mixin._on_realtime_telemetry_metrics(message)
+
+        # Should store empty list and still trigger hook
+        assert mocked_mixin._telemetry_metrics == []
+        mocked_mixin.run_hooks.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_concurrent_access_with_lock(self, mocked_mixin):
+        """Test that the lock protects concurrent access to telemetry metrics."""
+
+        # Track lock acquisition order
+        lock_acquired_order = []
+
+        async def acquire_lock_and_update(metrics_value, delay):
+            """Helper to simulate concurrent updates."""
+            async with mocked_mixin._telemetry_metrics_lock:
+                lock_acquired_order.append(metrics_value)
+                await asyncio.sleep(delay)
+                mocked_mixin._telemetry_metrics = [
+                    MetricResult(
+                        tag=f"metric_{metrics_value}",
+                        header=f"Metric {metrics_value}",
+                        unit="ms",
+                        avg=float(metrics_value),
+                    )
+                ]
+
+        # Start two concurrent operations
+        await asyncio.gather(
+            acquire_lock_and_update(1, 0.01), acquire_lock_and_update(2, 0.005)
+        )
+
+        # Both should have acquired the lock (order doesn't matter for this test)
+        assert len(lock_acquired_order) == 2
+        assert set(lock_acquired_order) == {1, 2}
+
+        # Final value should be from the last completed operation
+        assert len(mocked_mixin._telemetry_metrics) == 1
+
+    @pytest.mark.asyncio
+    async def test_multiple_metrics_handling(self, mocked_mixin):
+        """Test handling of message with multiple metrics."""
+        metrics = [
+            MetricResult(
+                tag=f"metric_{i}", header=f"Metric {i}", unit="ms", avg=float(i)
+            )
+            for i in range(10)
+        ]
+
+        message = RealtimeTelemetryMetricsMessage(
+            service_id="records_manager", metrics=metrics
+        )
+
+        await mocked_mixin._on_realtime_telemetry_metrics(message)
+
+        # All metrics should be stored
+        assert len(mocked_mixin._telemetry_metrics) == 10
+        assert mocked_mixin._telemetry_metrics == metrics
+
+    @pytest.mark.asyncio
+    async def test_integration_with_controller(self):
+        """Test that mixin integrates correctly with controller."""
+        service_config = ServiceConfig()
+        mock_controller = MagicMock()
+        mock_controller.some_method = MagicMock(return_value="test_value")
+
+        with patch(
+            "aiperf.common.mixins.message_bus_mixin.MessageBusClientMixin.__init__",
+            return_value=None,
+        ):
+            mixin = RealtimeTelemetryMetricsMixin(
+                service_config=service_config, controller=mock_controller
+            )
+
+        # Verify controller is accessible
+        assert mixin._controller == mock_controller
+        assert mixin._controller.some_method() == "test_value"
diff --git a/tests/ui/test_realtime_telemetry_dashboard.py b/tests/ui/test_realtime_telemetry_dashboard.py
