Merge pull request #4 from microsoft/nikhilc/supportS2S

Support S2S

Author: nikhilNava

libraries/microsoft-agents-a365-observability-core/microsoft_agents_a365/observability/core/exporters/agent365_exporter.py

@@ -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.")
@@ -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 -----------------
 
@@ -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:

tests/test_agent365_exporter.py

@@ -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()