Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ def __init__(
self,
token_resolver: Callable[[str, str], str | None],
cluster_category: str = "prod",
**kwargs: Any,
use_s2s_endpoint: bool = False,
):
if token_resolver is None:
raise ValueError("token_resolver must be provided.")
Expand All @@ -53,6 +53,7 @@ def __init__(
self._lock = threading.Lock()
self._token_resolver = token_resolver
self._cluster_category = cluster_category
self._use_s2s_endpoint = use_s2s_endpoint

# ------------- SpanExporter API -----------------

Expand All @@ -74,7 +75,12 @@ def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
# Resolve endpoint + token
discovery = PowerPlatformApiDiscovery(self._cluster_category)
endpoint = discovery.get_tenant_island_cluster_endpoint(tenant_id)
url = f"https://{endpoint}/maven/agent365/agents/{agent_id}/traces?api-version=1"
endpoint_path = (
f"/maven/agent365/service/agents/{agent_id}/traces"
if self._use_s2s_endpoint
else f"/maven/agent365/agents/{agent_id}/traces"
)
url = f"https://{endpoint}{endpoint_path}?api-version=1"

headers = {"content-type": "application/json"}
try:
Expand Down
74 changes: 74 additions & 0 deletions tests/test_agent365_exporter.py
Original file line number Diff line number Diff line change
Expand Up @@ -213,6 +213,80 @@ def test_partitioning_by_scope(self):
expected_scopes = [("scope.a", 2), ("scope.b", 1), ("scope.c", 1)]
self.assertEqual(scope_data, expected_scopes)

def test_s2s_endpoint_path_when_enabled(self):
"""Test 4: Test that S2S endpoint path is used when use_s2s_endpoint is True."""
# Arrange - Create exporter with S2S endpoint enabled
s2s_exporter = Agent365Exporter(
token_resolver=self.mock_token_resolver, cluster_category="test", use_s2s_endpoint=True
)

spans = [self._create_mock_span("s2s_span")]

# Mock the PowerPlatformApiDiscovery class
with patch(
"microsoft_agents_a365.observability.core.exporters.agent365_exporter.PowerPlatformApiDiscovery"
) as mock_discovery_class:
mock_discovery = Mock()
mock_discovery.get_tenant_island_cluster_endpoint.return_value = "test-endpoint.com"
mock_discovery_class.return_value = mock_discovery

# Mock the _post_with_retries method
with patch.object(s2s_exporter, "_post_with_retries", return_value=True) as mock_post:
# Act
result = s2s_exporter.export(spans)

# Assert
self.assertEqual(result, SpanExportResult.SUCCESS)
mock_post.assert_called_once()

# Verify the call arguments - should use S2S path
args, kwargs = mock_post.call_args
url, body, headers = args

self.assertIn("test-endpoint.com", url)
self.assertIn("/maven/agent365/service/agents/test-agent-456/traces", url)
self.assertNotIn("/maven/agent365/agents/test-agent-456/traces", url)
self.assertEqual(headers["authorization"], "Bearer test_token_123")
self.assertEqual(headers["content-type"], "application/json")

def test_default_endpoint_path_when_s2s_disabled(self):
"""Test 5: Test that default endpoint path is used when use_s2s_endpoint is False."""
# Arrange - Create exporter with S2S endpoint disabled (default behavior)
default_exporter = Agent365Exporter(
token_resolver=self.mock_token_resolver, cluster_category="test", use_s2s_endpoint=False
)

spans = [self._create_mock_span("default_span")]

# Mock the PowerPlatformApiDiscovery class
with patch(
"microsoft_agents_a365.observability.core.exporters.agent365_exporter.PowerPlatformApiDiscovery"
) as mock_discovery_class:
mock_discovery = Mock()
mock_discovery.get_tenant_island_cluster_endpoint.return_value = "test-endpoint.com"
mock_discovery_class.return_value = mock_discovery

# Mock the _post_with_retries method
with patch.object(
default_exporter, "_post_with_retries", return_value=True
) as mock_post:
# Act
result = default_exporter.export(spans)

# Assert
self.assertEqual(result, SpanExportResult.SUCCESS)
mock_post.assert_called_once()

# Verify the call arguments - should use default path
args, kwargs = mock_post.call_args
url, body, headers = args

self.assertIn("test-endpoint.com", url)
self.assertIn("/maven/agent365/agents/test-agent-456/traces", url)
self.assertNotIn("/maven/agent365/service/agents/test-agent-456/traces", url)
self.assertEqual(headers["authorization"], "Bearer test_token_123")
self.assertEqual(headers["content-type"], "application/json")


if __name__ == "__main__":
unittest.main()
Loading