Revert logic type to Callable[[TurnContext], Awaitable] to match Middleware Protocol; consolidate tests

Co-authored-by: nikhilNava <211831449+nikhilNava@users.noreply.github.com>

Author: Copilot

libraries/microsoft-agents-a365-observability-hosting/microsoft_agents_a365/observability/hosting/middleware/baggage_middleware.py

@@ -23,7 +23,7 @@ class BaggageMiddleware:
     async def on_turn(
         self,
         context: TurnContext,
-        logic: Callable[[], Awaitable],
+        logic: Callable[[TurnContext], Awaitable],
     ) -> None:
         activity = context.activity
         is_async_reply = (

libraries/microsoft-agents-a365-observability-hosting/microsoft_agents_a365/observability/hosting/middleware/output_logging_middleware.py

@@ -126,7 +126,7 @@ class OutputLoggingMiddleware:
     async def on_turn(
         self,
         context: TurnContext,
-        logic: Callable[[], Awaitable],
+        logic: Callable[[TurnContext], Awaitable],
     ) -> None:
         agent_details = _derive_agent_details(context)
         tenant_details = _derive_tenant_details(context)

tests/observability/hosting/middleware/test_baggage_middleware.py

@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import MagicMock
 
 import pytest
 from microsoft_agents.activity import (
@@ -95,15 +95,3 @@ async def logic():
     assert logic_called is True
     # Baggage should NOT be set because the middleware skipped it
     assert captured_caller_id is None
-
-
-@pytest.mark.asyncio
-async def test_baggage_middleware_calls_logic():
-    """BaggageMiddleware should always call the downstream logic."""
-    middleware = BaggageMiddleware()
-    ctx = _make_turn_context()
-
-    logic_mock = AsyncMock()
-    await middleware.on_turn(ctx, logic_mock)
-
-    logic_mock.assert_awaited_once()

tests/observability/hosting/middleware/test_observability_hosting_manager.py

@@ -24,73 +24,50 @@ def _reset_singleton():
     ObservabilityHostingManager.reset()
 
 
-def test_configure_returns_instance():
-    """configure() should return an ObservabilityHostingManager instance."""
-    middleware_set = MagicMock()
-    instance = ObservabilityHostingManager.configure(middleware_set, ObservabilityHostingOptions())
-    assert isinstance(instance, ObservabilityHostingManager)
-
-
 def test_configure_is_singleton():
-    """Subsequent configure() calls should return the same instance."""
+    """configure() should return an ObservabilityHostingManager and be a singleton."""
     middleware_set = MagicMock()
     options = ObservabilityHostingOptions()
     first = ObservabilityHostingManager.configure(middleware_set, options)
+    assert isinstance(first, ObservabilityHostingManager)
     second = ObservabilityHostingManager.configure(middleware_set, options)
     assert first is second
 
 
-def test_configure_registers_baggage_middleware_by_default():
-    """By default, BaggageMiddleware should be registered."""
-    middleware_set = MagicMock()
-    ObservabilityHostingManager.configure(middleware_set, ObservabilityHostingOptions())
-
-    # The middleware_set.use should have been called once (only BaggageMiddleware by default)
-    assert middleware_set.use.call_count == 1
-    registered = middleware_set.use.call_args_list[0][0][0]
-    assert isinstance(registered, BaggageMiddleware)
-
-
-def test_configure_registers_both_middlewares():
-    """When output logging is enabled, both middlewares should be registered."""
-    middleware_set = MagicMock()
-    options = ObservabilityHostingOptions(enable_baggage=True, enable_output_logging=True)
-    ObservabilityHostingManager.configure(middleware_set, options)
-
-    assert middleware_set.use.call_count == 2
-    registered_types = [c[0][0] for c in middleware_set.use.call_args_list]
-    assert isinstance(registered_types[0], BaggageMiddleware)
-    assert isinstance(registered_types[1], OutputLoggingMiddleware)
-
-
-def test_configure_disables_baggage():
-    """When baggage is disabled, only output logging should be registered (if enabled)."""
-    middleware_set = MagicMock()
-    options = ObservabilityHostingOptions(enable_baggage=False, enable_output_logging=True)
-    ObservabilityHostingManager.configure(middleware_set, options)
-
-    assert middleware_set.use.call_count == 1
-    registered = middleware_set.use.call_args_list[0][0][0]
-    assert isinstance(registered, OutputLoggingMiddleware)
-
-
-def test_configure_disables_all():
-    """When both are disabled, no middleware should be registered."""
+@pytest.mark.parametrize(
+    "enable_baggage,enable_output_logging,expected_types",
+    [
+        (True, False, [BaggageMiddleware]),
+        (True, True, [BaggageMiddleware, OutputLoggingMiddleware]),
+        (False, True, [OutputLoggingMiddleware]),
+        (False, False, []),
+    ],
+    ids=["default_baggage_only", "both_enabled", "output_only", "none"],
+)
+def test_configure_registers_expected_middlewares(
+    enable_baggage, enable_output_logging, expected_types
+):
+    """configure() should register the correct middlewares based on options."""
     middleware_set = MagicMock()
-    options = ObservabilityHostingOptions(enable_baggage=False, enable_output_logging=False)
+    options = ObservabilityHostingOptions(
+        enable_baggage=enable_baggage, enable_output_logging=enable_output_logging
+    )
     ObservabilityHostingManager.configure(middleware_set, options)
 
-    middleware_set.use.assert_not_called()
-
-
-def test_configure_raises_on_none_middleware_set():
-    """configure() should raise TypeError when middleware_set is None."""
-    with pytest.raises(TypeError, match="middleware_set must not be None"):
-        ObservabilityHostingManager.configure(None, ObservabilityHostingOptions())
+    assert middleware_set.use.call_count == len(expected_types)
+    for call, expected_type in zip(middleware_set.use.call_args_list, expected_types, strict=True):
+        assert isinstance(call[0][0], expected_type)
 
 
-def test_configure_raises_on_none_options():
-    """configure() should raise TypeError when options is None."""
-    middleware_set = MagicMock()
-    with pytest.raises(TypeError, match="options must not be None"):
-        ObservabilityHostingManager.configure(middleware_set, None)
+@pytest.mark.parametrize(
+    "middleware_set,options,match",
+    [
+        (None, ObservabilityHostingOptions(), "middleware_set must not be None"),
+        (MagicMock(), None, "options must not be None"),
+    ],
+    ids=["none_middleware_set", "none_options"],
+)
+def test_configure_raises_on_none(middleware_set, options, match):
+    """configure() should raise TypeError when required args are None."""
+    with pytest.raises(TypeError, match=match):
+        ObservabilityHostingManager.configure(middleware_set, options)

tests/observability/hosting/middleware/test_output_logging_middleware.py

@@ -136,7 +136,7 @@ async def test_send_handler_skips_non_message_activities():
 
 @pytest.mark.asyncio
 async def test_send_handler_creates_output_scope_for_messages():
-    """Send handler should create an OutputScope for message activities."""
+    """Send handler should create an OutputScope for message activities and dispose on success."""
     middleware = OutputLoggingMiddleware()
     ctx = _make_turn_context()
 
@@ -159,6 +159,7 @@ async def test_send_handler_creates_output_scope_for_messages():
         mock_output_scope_cls.start.assert_called_once()
         send_next.assert_awaited_once()
         mock_scope.dispose.assert_called_once()
+        mock_scope.record_error.assert_not_called()
 
 
 @pytest.mark.asyncio
@@ -216,29 +217,3 @@ async def test_send_handler_rethrows_errors():
 
         mock_scope.record_error.assert_called_once_with(send_error)
         mock_scope.dispose.assert_called_once()
-
-
-@pytest.mark.asyncio
-async def test_send_handler_disposes_scope_on_success():
-    """Send handler should dispose the OutputScope even when send_next succeeds."""
-    middleware = OutputLoggingMiddleware()
-    ctx = _make_turn_context()
-
-    await middleware.on_turn(ctx, AsyncMock())
-
-    handler = ctx._on_send_activities[-1]
-
-    activities = [Activity(type="message", text="Reply")]
-    send_next = AsyncMock()
-
-    with patch(
-        "microsoft_agents_a365.observability.hosting.middleware"
-        ".output_logging_middleware.OutputScope"
-    ) as mock_output_scope_cls:
-        mock_scope = MagicMock()
-        mock_output_scope_cls.start.return_value = mock_scope
-
-        await handler(ctx, activities, send_next)
-
-        mock_scope.dispose.assert_called_once()
-        mock_scope.record_error.assert_not_called()