mirror of
https://github.com/langgenius/dify.git
synced 2026-02-10 04:54:30 +08:00
4a9b74f86b
**Problem:**
The telemetry system had unnecessary abstraction layers and bad practices
from the last 3 commits introducing the gateway implementation:
- TelemetryFacade class wrapper around emit() function
- String literals instead of SignalType enum
- Dictionary mapping enum → string instead of enum → enum
- Unnecessary ENTERPRISE_TELEMETRY_GATEWAY_ENABLED feature flag
- Duplicate guard checks scattered across files
- Non-thread-safe TelemetryGateway singleton pattern
- Missing guard in ops_trace_task.py causing RuntimeError spam
**Solution:**
1. Deleted TelemetryFacade - replaced with thin emit() function in core/telemetry/__init__.py
2. Added SignalType enum ('trace' | 'metric_log') to enterprise/telemetry/contracts.py
3. Replaced CASE_TO_TRACE_TASK_NAME dict with CASE_TO_TRACE_TASK: dict[TelemetryCase, TraceTaskName]
4. Deleted is_gateway_enabled() and _emit_legacy() - using existing ENTERPRISE_ENABLED + ENTERPRISE_TELEMETRY_ENABLED instead
5. Extracted _should_drop_ee_only_event() helper to eliminate duplicate checks
6. Moved TelemetryGateway singleton to ext_enterprise_telemetry.py:
- Init once in init_app() for thread-safety
- Access via get_gateway() function
7. Re-added guard to ops_trace_task.py to prevent RuntimeError when EE=OFF but CE tracing enabled
8. Updated 11 caller files to import 'emit as telemetry_emit' instead of 'TelemetryFacade'
**Result:**
- 322 net lines deleted (533 removed, 211 added)
- All 91 tests pass
- Thread-safe singleton pattern
- Cleaner API surface: from TelemetryFacade.emit() to telemetry_emit()
- Proper enum usage throughout
- No RuntimeError spam in EE=OFF + CE=ON scenario
70 lines
2.2 KiB
Python
70 lines
2.2 KiB
Python
"""Unit tests for enterprise telemetry Celery task."""
|
|
|
|
import json
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from enterprise.telemetry.contracts import TelemetryCase, TelemetryEnvelope
|
|
from tasks.enterprise_telemetry_task import process_enterprise_telemetry
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_envelope_json():
|
|
envelope = TelemetryEnvelope(
|
|
case=TelemetryCase.APP_CREATED,
|
|
tenant_id="test-tenant",
|
|
event_id="test-event-123",
|
|
payload={"app_id": "app-123"},
|
|
)
|
|
return envelope.model_dump_json()
|
|
|
|
|
|
def test_process_enterprise_telemetry_success(sample_envelope_json):
|
|
with patch("tasks.enterprise_telemetry_task.EnterpriseMetricHandler") as mock_handler_class:
|
|
mock_handler = MagicMock()
|
|
mock_handler_class.return_value = mock_handler
|
|
|
|
process_enterprise_telemetry(sample_envelope_json)
|
|
|
|
mock_handler.handle.assert_called_once()
|
|
call_args = mock_handler.handle.call_args[0][0]
|
|
assert isinstance(call_args, TelemetryEnvelope)
|
|
assert call_args.case == TelemetryCase.APP_CREATED
|
|
assert call_args.tenant_id == "test-tenant"
|
|
assert call_args.event_id == "test-event-123"
|
|
|
|
|
|
def test_process_enterprise_telemetry_invalid_json(caplog):
|
|
invalid_json = "not valid json"
|
|
|
|
process_enterprise_telemetry(invalid_json)
|
|
|
|
assert "Failed to process enterprise telemetry envelope" in caplog.text
|
|
|
|
|
|
def test_process_enterprise_telemetry_handler_exception(sample_envelope_json, caplog):
|
|
with patch("tasks.enterprise_telemetry_task.EnterpriseMetricHandler") as mock_handler_class:
|
|
mock_handler = MagicMock()
|
|
mock_handler.handle.side_effect = Exception("Handler error")
|
|
mock_handler_class.return_value = mock_handler
|
|
|
|
process_enterprise_telemetry(sample_envelope_json)
|
|
|
|
assert "Failed to process enterprise telemetry envelope" in caplog.text
|
|
|
|
|
|
def test_process_enterprise_telemetry_validation_error(caplog):
|
|
invalid_envelope = json.dumps(
|
|
{
|
|
"case": "INVALID_CASE",
|
|
"tenant_id": "test-tenant",
|
|
"event_id": "test-event",
|
|
"payload": {},
|
|
}
|
|
)
|
|
|
|
process_enterprise_telemetry(invalid_envelope)
|
|
|
|
assert "Failed to process enterprise telemetry envelope" in caplog.text
|