From 164228117659e40eda01cb46d818367229a500da Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Wed, 30 Sep 2026 06:21:01 +0000 Subject: [PATCH] test: exercise real telemetry queues and trace managers (#43162) --- .../core/ops/test_trace_queue_manager.py | 90 ++++++------------- .../unit_tests/core/telemetry/test_facade.py | 71 +++++---------- 2 files changed, 52 insertions(+), 109 deletions(-) diff --git a/api/tests/unit_tests/core/ops/test_trace_queue_manager.py b/api/tests/unit_tests/core/ops/test_trace_queue_manager.py index 13cf01651e3..045bf02235d 100644 --- a/api/tests/unit_tests/core/ops/test_trace_queue_manager.py +++ b/api/tests/unit_tests/core/ops/test_trace_queue_manager.py @@ -11,58 +11,19 @@ When BOTH are false, tasks are silently dropped (correct behavior). """ import queue -import sys -import types from unittest.mock import MagicMock, patch import pytest @pytest.fixture -def trace_queue_manager_and_task(monkeypatch: pytest.MonkeyPatch): - """Fixture to provide TraceQueueManager and TraceTask with delayed imports.""" - module_name = "core.ops.ops_trace_manager" - if module_name not in sys.modules: - ops_stub = types.ModuleType(module_name) - - class StubTraceTask: - def __init__(self, trace_type): - self.trace_type = trace_type - self.app_id = None - - class StubTraceQueueManager: - def __init__(self, app_id=None): - self.app_id = app_id - from core.telemetry.gateway import is_enterprise_telemetry_enabled - - self._enterprise_telemetry_enabled = is_enterprise_telemetry_enabled() - self.trace_instance = StubOpsTraceManager.get_ops_trace_instance(app_id) - - def add_trace_task(self, trace_task): - if self._enterprise_telemetry_enabled or self.trace_instance: - trace_task.app_id = self.app_id - from core.ops.ops_trace_manager import trace_manager_queue - - trace_manager_queue.put(trace_task) - - class StubOpsTraceManager: - @staticmethod - def get_ops_trace_instance(app_id): - return None - - ops_stub.TraceQueueManager = StubTraceQueueManager - ops_stub.TraceTask = StubTraceTask - ops_stub.OpsTraceManager = StubOpsTraceManager - ops_stub.trace_manager_queue = MagicMock(spec=queue.Queue) - monkeypatch.setitem(sys.modules, module_name, ops_stub) - +def trace_queue_manager_and_task(): + """Exercise the production guard without starting background dispatch.""" from core.ops.entities.trace_entity import TraceTaskName + from core.ops.ops_trace_manager import TraceQueueManager, TraceTask - ops_module = __import__(module_name, fromlist=["TraceQueueManager", "TraceTask"]) - TraceQueueManager = ops_module.TraceQueueManager - TraceTask = ops_module.TraceTask - - return TraceQueueManager, TraceTask, TraceTaskName + with patch.object(TraceQueueManager, "start_timer"): + yield TraceQueueManager, TraceTask, TraceTaskName class TestTraceQueueManagerTelemetryGuard: @@ -76,19 +37,19 @@ class TestTraceQueueManagerTelemetryGuard: """ TraceQueueManager, TraceTask, TraceTaskName = trace_queue_manager_and_task - mock_queue = MagicMock(spec=queue.Queue) + trace_queue = queue.Queue() trace_task = TraceTask(trace_type=TraceTaskName.WORKFLOW_TRACE) with ( patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False), patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None), - patch("core.ops.ops_trace_manager.trace_manager_queue", mock_queue), + patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), ): manager = TraceQueueManager(app_id="test-app-id") manager.add_trace_task(trace_task) - mock_queue.put.assert_not_called() + assert trace_queue.empty() def test_task_enqueued_when_telemetry_enabled(self, trace_queue_manager_and_task): """Verify task IS enqueued when enterprise telemetry is enabled. @@ -98,20 +59,21 @@ class TestTraceQueueManagerTelemetryGuard: """ TraceQueueManager, TraceTask, TraceTaskName = trace_queue_manager_and_task - mock_queue = MagicMock(spec=queue.Queue) + trace_queue = queue.Queue() trace_task = TraceTask(trace_type=TraceTaskName.WORKFLOW_TRACE) with ( patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True), patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None), - patch("core.ops.ops_trace_manager.trace_manager_queue", mock_queue), + patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), ): manager = TraceQueueManager(app_id="test-app-id") manager.add_trace_task(trace_task) - mock_queue.put.assert_called_once() - called_task = mock_queue.put.call_args[0][0] + called_task = trace_queue.get_nowait() + assert called_task is trace_task + assert trace_queue.empty() assert called_task.app_id == "test-app-id" def test_task_enqueued_when_trace_instance_configured(self, trace_queue_manager_and_task): @@ -122,7 +84,7 @@ class TestTraceQueueManagerTelemetryGuard: """ TraceQueueManager, TraceTask, TraceTaskName = trace_queue_manager_and_task - mock_queue = MagicMock(spec=queue.Queue) + trace_queue = queue.Queue() mock_trace_instance = MagicMock() @@ -133,13 +95,14 @@ class TestTraceQueueManagerTelemetryGuard: patch( "core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=mock_trace_instance ), - patch("core.ops.ops_trace_manager.trace_manager_queue", mock_queue), + patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), ): manager = TraceQueueManager(app_id="test-app-id") manager.add_trace_task(trace_task) - mock_queue.put.assert_called_once() - called_task = mock_queue.put.call_args[0][0] + called_task = trace_queue.get_nowait() + assert called_task is trace_task + assert trace_queue.empty() assert called_task.app_id == "test-app-id" def test_task_enqueued_when_both_telemetry_and_trace_instance_enabled(self, trace_queue_manager_and_task): @@ -150,7 +113,7 @@ class TestTraceQueueManagerTelemetryGuard: """ TraceQueueManager, TraceTask, TraceTaskName = trace_queue_manager_and_task - mock_queue = MagicMock(spec=queue.Queue) + trace_queue = queue.Queue() mock_trace_instance = MagicMock() @@ -161,13 +124,14 @@ class TestTraceQueueManagerTelemetryGuard: patch( "core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=mock_trace_instance ), - patch("core.ops.ops_trace_manager.trace_manager_queue", mock_queue), + patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), ): manager = TraceQueueManager(app_id="test-app-id") manager.add_trace_task(trace_task) - mock_queue.put.assert_called_once() - called_task = mock_queue.put.call_args[0][0] + called_task = trace_queue.get_nowait() + assert called_task is trace_task + assert trace_queue.empty() assert called_task.app_id == "test-app-id" def test_app_id_set_before_enqueue(self, trace_queue_manager_and_task): @@ -178,17 +142,19 @@ class TestTraceQueueManagerTelemetryGuard: """ TraceQueueManager, TraceTask, TraceTaskName = trace_queue_manager_and_task - mock_queue = MagicMock(spec=queue.Queue) + trace_queue = queue.Queue() trace_task = TraceTask(trace_type=TraceTaskName.WORKFLOW_TRACE) with ( patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True), patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None), - patch("core.ops.ops_trace_manager.trace_manager_queue", mock_queue), + patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), ): manager = TraceQueueManager(app_id="expected-app-id") manager.add_trace_task(trace_task) - called_task = mock_queue.put.call_args[0][0] + called_task = trace_queue.get_nowait() + assert called_task is trace_task + assert trace_queue.empty() assert called_task.app_id == "expected-app-id" diff --git a/api/tests/unit_tests/core/telemetry/test_facade.py b/api/tests/unit_tests/core/telemetry/test_facade.py index 6606b56f911..899d4a537dd 100644 --- a/api/tests/unit_tests/core/telemetry/test_facade.py +++ b/api/tests/unit_tests/core/telemetry/test_facade.py @@ -3,59 +3,32 @@ from __future__ import annotations import queue -import sys -import types from unittest.mock import MagicMock, patch import pytest from core.ops.entities.trace_entity import TraceTaskName +from core.ops.ops_trace_manager import TraceQueueManager, TraceTask +from core.telemetry import emit from core.telemetry.events import DraftNodeExecutionTraceEvent, TelemetryContext @pytest.fixture -def telemetry_test_setup(monkeypatch: pytest.MonkeyPatch): - module_name = "core.ops.ops_trace_manager" - ops_stub = types.ModuleType(module_name) - - class StubTraceTask: - def __init__(self, trace_type, **kwargs): - self.trace_type = trace_type - self.app_id = None - self.kwargs = kwargs - - class StubTraceQueueManager: - def __init__(self, app_id=None, user_id=None): - self.app_id = app_id - self.user_id = user_id - self.trace_instance = StubOpsTraceManager.get_ops_trace_instance(app_id) - - def add_trace_task(self, trace_task): - trace_task.app_id = self.app_id - from core.ops.ops_trace_manager import trace_manager_queue - - trace_manager_queue.put(trace_task) - - class StubOpsTraceManager: - @staticmethod - def get_ops_trace_instance(app_id): - return None - - ops_stub.TraceQueueManager = StubTraceQueueManager - ops_stub.TraceTask = StubTraceTask - ops_stub.OpsTraceManager = StubOpsTraceManager - ops_stub.trace_manager_queue = MagicMock(spec=queue.Queue) - monkeypatch.setitem(sys.modules, module_name, ops_stub) - - from core.telemetry import emit - - return emit, ops_stub.trace_manager_queue +def telemetry_test_setup(): + trace_queue: queue.Queue[TraceTask] = queue.Queue() + with ( + patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None), + patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), + patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False), + patch.object(TraceQueueManager, "start_timer"), + ): + yield emit, trace_queue class TestTelemetryEmit: @patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True) def test_emit_enterprise_trace_creates_trace_task(self, mock_ee, telemetry_test_setup): - emit_fn, mock_queue = telemetry_test_setup + emit_fn, trace_queue = telemetry_test_setup event = DraftNodeExecutionTraceEvent( context=TelemetryContext( @@ -68,12 +41,14 @@ class TestTelemetryEmit: emit_fn(event) - mock_queue.put.assert_called_once() - called_task = mock_queue.put.call_args[0][0] + called_task = trace_queue.get_nowait() + assert trace_queue.empty() + assert called_task.app_id == "test-app" + assert called_task.user_id == "test-user" assert called_task.trace_type == TraceTaskName.DRAFT_NODE_EXECUTION_TRACE def test_emit_enterprise_only_trace_dropped_when_ee_disabled(self, telemetry_test_setup): - emit_fn, mock_queue = telemetry_test_setup + emit_fn, trace_queue = telemetry_test_setup event = DraftNodeExecutionTraceEvent( context=TelemetryContext( @@ -86,11 +61,11 @@ class TestTelemetryEmit: emit_fn(event) - mock_queue.put.assert_not_called() + assert trace_queue.empty() @patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True) def test_emit_passes_name_directly_to_trace_task(self, mock_ee, telemetry_test_setup): - emit_fn, mock_queue = telemetry_test_setup + emit_fn, trace_queue = telemetry_test_setup event = DraftNodeExecutionTraceEvent( context=TelemetryContext( @@ -103,14 +78,16 @@ class TestTelemetryEmit: emit_fn(event) - mock_queue.put.assert_called_once() - called_task = mock_queue.put.call_args[0][0] + called_task = trace_queue.get_nowait() + assert trace_queue.empty() + assert called_task.app_id == "test-app" + assert called_task.user_id == "test-user" assert called_task.trace_type == TraceTaskName.DRAFT_NODE_EXECUTION_TRACE assert isinstance(called_task.trace_type, TraceTaskName) @patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True) def test_emit_with_provided_trace_manager(self, mock_ee, telemetry_test_setup): - emit_fn, mock_queue = telemetry_test_setup + emit_fn, trace_queue = telemetry_test_setup mock_trace_manager = MagicMock() mock_trace_manager.add_trace_task = MagicMock()