add additional test case

Author: liustve

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_starlette_patches.py

@@ -4,7 +4,7 @@
 from logging import Logger, getLogger
 from typing import Collection
 
-from amazon.opentelemetry.distro._utils import AGENT_OBSERVABILITY_ENABLED
+from amazon.opentelemetry.distro._utils import is_agent_observability_enabled
 
 _logger: Logger = getLogger(__name__)
 
@@ -31,10 +31,10 @@ def patched_instrumentation_dependencies(self) -> Collection[str]:
         # Apply the patch
         StarletteInstrumentor.instrumentation_dependencies = patched_instrumentation_dependencies
 
-        # Patch to exclude http receive/send ASGI event spans, this Middleware instrumentation is injected
-        # internally by Starlette Instrumentor, see:
+        # Patch to exclude http receive/send ASGI event spans from Bedrock AgentCore,
+        # this Middleware instrumentation is injected internally by Starlette Instrumentor, see:
         # https://github.com/open-telemetry/opentelemetry-python-contrib/blob/51da0a766e5d3cbc746189e10c9573163198cfcd/instrumentation/opentelemetry-instrumentation-asgi/src/opentelemetry/instrumentation/asgi/__init__.py#L573
-        if AGENT_OBSERVABILITY_ENABLED:
+        if is_agent_observability_enabled():
             original_init = OpenTelemetryMiddleware.__init__
 
             def patched_init(self, app, **kwargs):

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_starlette_patches.py

@@ -1,71 +1,74 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
-from unittest import TestCase
+import unittest
 from unittest.mock import MagicMock, patch
 
 from amazon.opentelemetry.distro.patches._starlette_patches import _apply_starlette_instrumentation_patches
 
 
-class TestStarlettePatch(TestCase):
+class TestStarlettePatch(unittest.TestCase):
     """Test the Starlette instrumentation patches."""
 
-    @patch("amazon.opentelemetry.distro.patches._starlette_patches.AGENT_OBSERVABILITY_ENABLED", True)
     @patch("amazon.opentelemetry.distro.patches._starlette_patches._logger")
     def test_starlette_patch_applied_successfully(self, mock_logger):
         """Test that the Starlette instrumentation patch is applied successfully."""
-        # Create a mock StarletteInstrumentor class
-        mock_instrumentor_class = MagicMock()
-        mock_instrumentor_class.__name__ = "StarletteInstrumentor"
-
-        class MockMiddleware:
-            def __init__(self, app, **kwargs):
-                pass
-
-        mock_middleware_class = MockMiddleware
-        original_init = mock_middleware_class.__init__
-
-        # Create mock modules
-        mock_starlette_module = MagicMock()
-        mock_starlette_module.StarletteInstrumentor = mock_instrumentor_class
-
-        mock_asgi_module = MagicMock()
-        mock_asgi_module.OpenTelemetryMiddleware = mock_middleware_class
-
-        # Mock the imports
-        with patch.dict(
-            "sys.modules",
-            {
-                "opentelemetry.instrumentation.starlette": mock_starlette_module,
-                "opentelemetry.instrumentation.asgi": mock_asgi_module,
-            },
-        ):
-            # Apply the patch
-            _apply_starlette_instrumentation_patches()
-
-            # Verify the instrumentation_dependencies method was replaced
-            self.assertTrue(hasattr(mock_instrumentor_class, "instrumentation_dependencies"))
-
-            # Test the patched method returns the expected value
-            mock_instance = MagicMock()
-            result = mock_instrumentor_class.instrumentation_dependencies(mock_instance)
-            self.assertEqual(result, ("starlette >= 0.13",))
-
-            self.assertNotEqual(mock_middleware_class.__init__, original_init)
-
-            # Test middleware patching sets exclude flags
-            mock_middleware_instance = MagicMock()
-            mock_middleware_instance.exclude_receive_span = False
-            mock_middleware_instance.exclude_send_span = False
-
-            mock_middleware_class.__init__(mock_middleware_instance, "app")
-
-            self.assertTrue(mock_middleware_instance.exclude_receive_span)
-            self.assertTrue(mock_middleware_instance.exclude_send_span)
-
-            # Verify logging
-            mock_logger.debug.assert_called_once_with(
-                "Successfully patched Starlette instrumentation_dependencies method"
-            )
+        for agent_enabled in [True, False]:
+            with self.subTest(agent_enabled=agent_enabled):
+                with patch.dict("os.environ", {"AGENT_OBSERVABILITY_ENABLED": "true" if agent_enabled else "false"}):
+                    # Create a mock StarletteInstrumentor class
+                    mock_instrumentor_class = MagicMock()
+                    mock_instrumentor_class.__name__ = "StarletteInstrumentor"
+
+                    def create_middleware_class():
+                        class MockMiddleware:
+                            def __init__(self, app, **kwargs):
+                                pass
+
+                        return MockMiddleware
+
+                    mock_middleware_class = create_middleware_class()
+
+                    mock_starlette_module = MagicMock()
+                    mock_starlette_module.StarletteInstrumentor = mock_instrumentor_class
+
+                    mock_asgi_module = MagicMock()
+                    mock_asgi_module.OpenTelemetryMiddleware = mock_middleware_class
+
+                    with patch.dict(
+                        "sys.modules",
+                        {
+                            "opentelemetry.instrumentation.starlette": mock_starlette_module,
+                            "opentelemetry.instrumentation.asgi": mock_asgi_module,
+                        },
+                    ):
+                        # Apply the patch
+                        _apply_starlette_instrumentation_patches()
+
+                        # Verify the instrumentation_dependencies method was replaced
+                        self.assertTrue(hasattr(mock_instrumentor_class, "instrumentation_dependencies"))
+
+                        # Test the patched method returns the expected value
+                        mock_instance = MagicMock()
+                        result = mock_instrumentor_class.instrumentation_dependencies(mock_instance)
+                        self.assertEqual(result, ("starlette >= 0.13",))
+
+                        mock_middleware_instance = MagicMock()
+                        mock_middleware_instance.exclude_receive_span = False
+                        mock_middleware_instance.exclude_send_span = False
+                        mock_middleware_class.__init__(mock_middleware_instance, "app")
+
+                        # Test middleware patching sets exclude flags
+                        if agent_enabled:
+                            self.assertTrue(mock_middleware_instance.exclude_receive_span)
+                            self.assertTrue(mock_middleware_instance.exclude_send_span)
+                        else:
+                            self.assertFalse(mock_middleware_instance.exclude_receive_span)
+                            self.assertFalse(mock_middleware_instance.exclude_send_span)
+
+                        # Verify logging
+                        mock_logger.debug.assert_called_with(
+                            "Successfully patched Starlette instrumentation_dependencies method"
+                        )
 
     @patch("amazon.opentelemetry.distro.patches._starlette_patches._logger")
     def test_starlette_patch_handles_import_error(self, mock_logger):
@@ -108,7 +111,7 @@ class MockStarletteInstrumentor(metaclass=ErrorMeta):
             args = mock_logger.warning.call_args[0]
             self.assertIn("Failed to apply Starlette instrumentation patches", args[0])
 
-    def test_starlette_patch_logs_failure_with_no_logger_patch(self):  # pylint: disable=no-self-use
+    def test_starlette_patch_logs_failure_with_no_logger_patch(self):
         """Test that the patch handles exceptions gracefully without logger mock."""
         # Mock the import to fail
         with patch.dict("sys.modules", {"opentelemetry.instrumentation.starlette": None}):