mirror of
https://github.com/langgenius/dify.git
synced 2026-10-04 13:47:55 +08:00
test: exercise real telemetry queues and trace managers (#43162)
This commit is contained in:
Asuka Minato authored and 
GitHub committed
1 parent
8edc40e2f1
commit
1642281176
2 files changed
+52
-109
No files matched your search
@@ -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"
|
||||
@@ -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()
|
||||
|
||||
Reference in new issue
Block a user