mirror of
https://github.com/langgenius/dify.git
synced 2026-02-24 18:05:11 +00:00
feat: inject traceparent in enterprise api (#30895)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
@@ -1,9 +1,14 @@
|
|||||||
|
import logging
|
||||||
import os
|
import os
|
||||||
from collections.abc import Mapping
|
from collections.abc import Mapping
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
import httpx
|
import httpx
|
||||||
|
|
||||||
|
from core.helper.trace_id_helper import generate_traceparent_header
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
class BaseRequest:
|
class BaseRequest:
|
||||||
proxies: Mapping[str, str] | None = {
|
proxies: Mapping[str, str] | None = {
|
||||||
@@ -38,6 +43,15 @@ class BaseRequest:
|
|||||||
headers = {"Content-Type": "application/json", cls.secret_key_header: cls.secret_key}
|
headers = {"Content-Type": "application/json", cls.secret_key_header: cls.secret_key}
|
||||||
url = f"{cls.base_url}{endpoint}"
|
url = f"{cls.base_url}{endpoint}"
|
||||||
mounts = cls._build_mounts()
|
mounts = cls._build_mounts()
|
||||||
|
|
||||||
|
try:
|
||||||
|
# ensure traceparent even when OTEL is disabled
|
||||||
|
traceparent = generate_traceparent_header()
|
||||||
|
if traceparent:
|
||||||
|
headers["traceparent"] = traceparent
|
||||||
|
except Exception:
|
||||||
|
logger.debug("Failed to generate traceparent header", exc_info=True)
|
||||||
|
|
||||||
with httpx.Client(mounts=mounts) as client:
|
with httpx.Client(mounts=mounts) as client:
|
||||||
response = client.request(method, url, json=json, params=params, headers=headers)
|
response = client.request(method, url, json=json, params=params, headers=headers)
|
||||||
return response.json()
|
return response.json()
|
||||||
|
|||||||
@@ -0,0 +1,59 @@
|
|||||||
|
"""Unit tests for traceparent header propagation in EnterpriseRequest.
|
||||||
|
|
||||||
|
This test module verifies that the W3C traceparent header is properly
|
||||||
|
generated and included in HTTP requests made by EnterpriseRequest.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from services.enterprise.base import EnterpriseRequest
|
||||||
|
|
||||||
|
|
||||||
|
class TestTraceparentPropagation:
|
||||||
|
"""Unit tests for traceparent header propagation."""
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_enterprise_config(self):
|
||||||
|
"""Mock EnterpriseRequest configuration."""
|
||||||
|
with (
|
||||||
|
patch.object(EnterpriseRequest, "base_url", "https://enterprise-api.example.com"),
|
||||||
|
patch.object(EnterpriseRequest, "secret_key", "test-secret-key"),
|
||||||
|
patch.object(EnterpriseRequest, "secret_key_header", "Enterprise-Api-Secret-Key"),
|
||||||
|
):
|
||||||
|
yield
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_httpx_client(self):
|
||||||
|
"""Mock httpx.Client for testing."""
|
||||||
|
with patch("services.enterprise.base.httpx.Client") as mock_client_class:
|
||||||
|
mock_client = MagicMock()
|
||||||
|
mock_client_class.return_value.__enter__.return_value = mock_client
|
||||||
|
mock_client_class.return_value.__exit__.return_value = None
|
||||||
|
|
||||||
|
# Setup default response
|
||||||
|
mock_response = MagicMock()
|
||||||
|
mock_response.json.return_value = {"result": "success"}
|
||||||
|
mock_client.request.return_value = mock_response
|
||||||
|
|
||||||
|
yield mock_client
|
||||||
|
|
||||||
|
def test_traceparent_header_included_when_generated(self, mock_enterprise_config, mock_httpx_client):
|
||||||
|
"""Test that traceparent header is included when successfully generated."""
|
||||||
|
# Arrange
|
||||||
|
expected_traceparent = "00-5b8aa5a2d2c872e8321cf37308d69df2-051581bf3bb55c45-01"
|
||||||
|
|
||||||
|
with patch("services.enterprise.base.generate_traceparent_header", return_value=expected_traceparent):
|
||||||
|
# Act
|
||||||
|
EnterpriseRequest.send_request("GET", "/test")
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
mock_httpx_client.request.assert_called_once()
|
||||||
|
call_args = mock_httpx_client.request.call_args
|
||||||
|
headers = call_args[1]["headers"]
|
||||||
|
|
||||||
|
assert "traceparent" in headers
|
||||||
|
assert headers["traceparent"] == expected_traceparent
|
||||||
|
assert headers["Content-Type"] == "application/json"
|
||||||
|
assert headers["Enterprise-Api-Secret-Key"] == "test-secret-key"
|
||||||
Reference in New Issue
Block a user