fix: update test methods to use new workflow run interface and improve context handling

Author: Dhanushree-Microsoft

src/processor/src/tests/unit/libs/agent_framework/test_agent_builder.py

@@ -157,15 +157,17 @@ def test_build_passes_all_state_to_chat_agent(self):
                 .build()
             )
         assert agent is mock_chat.return_value
+        args = mock_chat.call_args.args
         kwargs = mock_chat.call_args.kwargs
-        assert kwargs["chat_client"] is chat_client
+        assert args[0] is chat_client
         assert kwargs["instructions"] == "inst"
         assert kwargs["id"] == "id1"
         assert kwargs["name"] == "name1"
         assert kwargs["description"] == "desc1"
-        assert kwargs["temperature"] == 0.3
-        assert kwargs["max_tokens"] == 100
-        assert kwargs["tool_choice"] == "auto"
+        opts = kwargs["default_options"]
+        assert opts["temperature"] == 0.3
+        assert opts["max_tokens"] == 100
+        assert opts["tool_choice"] == "auto"
         assert kwargs["extra"] == 42
 
 
@@ -180,11 +182,13 @@ def test_create_agent_invokes_chat_agent(self):
                 temperature=0.4,
             )
         assert agent is mock_chat.return_value
+        args = mock_chat.call_args.args
         kwargs = mock_chat.call_args.kwargs
-        assert kwargs["chat_client"] is chat_client
+        assert args[0] is chat_client
         assert kwargs["instructions"] == "i"
         assert kwargs["name"] == "n"
-        assert kwargs["temperature"] == 0.4
+        opts = kwargs["default_options"]
+        assert opts["temperature"] == 0.4
 
     def test_create_agent_by_agentinfo_uses_helper_and_creates_client(self):
         # Build a fake AgentInfo with the minimum surface used by the method
@@ -215,12 +219,14 @@ def test_create_agent_by_agentinfo_uses_helper_and_creates_client(self):
         assert agent is mock_chat.return_value
         helper.settings.get_service_config.assert_called_once_with("default")
         helper.create_client.assert_called_once()
+        args = mock_chat.call_args.args
         ck = mock_chat.call_args.kwargs
-        assert ck["chat_client"] == "client-instance"
+        assert args[0] == "client-instance"
         assert ck["instructions"] == "instr"
         assert ck["name"] == "A"
         assert ck["description"] == "D"
-        assert ck["temperature"] == 0.2
+        opts = ck["default_options"]
+        assert opts["temperature"] == 0.2
 
     def test_create_agent_by_agentinfo_falls_back_to_system_prompt(self):
         helper = MagicMock()

src/processor/src/tests/unit/libs/agent_framework/test_groupchat_orchestrator_internals.py

@@ -743,16 +743,14 @@ def test_build_groupchat_invokes_builder(self):
         })
         with patch("libs.agent_framework.groupchat_orchestrator.GroupChatBuilder") as MockBuilder:
             built = MagicMock()
-            built.set_manager.return_value = built
-            built.participants.return_value = built
             built.build.return_value = "wf"
             MockBuilder.return_value = built
             wf = _run(orch._build_groupchat())
         assert wf == "wf"
         # ResultGenerator excluded from participants
-        kwargs = built.participants.call_args.args[0]
-        assert "arch" in kwargs
-        assert "rg" not in kwargs
+        kwargs = MockBuilder.call_args.kwargs
+        assert "arch" in kwargs["participants"]
+        assert "rg" not in kwargs["participants"]
 
 
 # -----------------------------------------------------------------------------

src/processor/src/tests/unit/libs/agent_framework/test_shared_memory_context_provider.py

@@ -61,8 +61,36 @@ def _make_provider(store=None):
     ), store
 
 
+def _make_context(messages=None, response_text=None):
+    """Create a mock SessionContext for before_run/after_run calls."""
+    ctx = MagicMock()
+    ctx.get_messages = MagicMock(return_value=messages or [])
+    ctx.extend_instructions = MagicMock()
+    if response_text is not None:
+        ctx.response = MagicMock()
+        ctx.response.text = response_text
+    else:
+        ctx.response = None
+    return ctx
+
+
+async def _call_before_run(provider, messages):
+    """Helper to call before_run and return the instructions that were injected."""
+    ctx = _make_context(messages=messages)
+    await provider.before_run(agent=MagicMock(), session=MagicMock(), context=ctx, state={})
+    if ctx.extend_instructions.called:
+        return ctx.extend_instructions.call_args[0][1]  # second positional arg = instructions
+    return None
+
+
+async def _call_after_run(provider, response_text):
+    """Helper to call after_run with a response."""
+    ctx = _make_context(response_text=response_text)
+    await provider.after_run(agent=MagicMock(), session=MagicMock(), context=ctx, state={})
+
+
 # ---------------------------------------------------------------------------
-# invoking() — Pre-LLM memory injection
+# before_run() — Pre-LLM memory injection
 # ---------------------------------------------------------------------------
 
 
@@ -75,11 +103,11 @@ async def _run():
         ]
         messages = [_make_chat_message("How should we handle storage configuration?")]
 
-        context = await provider.invoking(messages)
+        instructions = await _call_before_run(provider, messages)
 
-        assert context.instructions is not None
-        assert "GKE Filestore CSI" in context.instructions
-        assert "Azure Files for AKS" in context.instructions
+        assert instructions is not None
+        assert "GKE Filestore CSI" in instructions
+        assert "Azure Files for AKS" in instructions
         store.search.assert_called_once()
 
     asyncio.run(_run())
@@ -88,9 +116,8 @@ async def _run():
 def test_invoking_empty_messages_returns_empty():
     async def _run():
         provider, _ = _make_provider()
-        context = await provider.invoking([])
-        assert context.instructions is None
-        assert getattr(context, "messages", []) == []
+        instructions = await _call_before_run(provider, [])
+        assert instructions is None
 
     asyncio.run(_run())
 
@@ -101,8 +128,8 @@ async def _run():
         store.search.return_value = []
         messages = [_make_chat_message("What is the overall migration plan for AKS?")]
 
-        context = await provider.invoking(messages)
-        assert context.instructions is None
+        instructions = await _call_before_run(provider, messages)
+        assert instructions is None
 
     asyncio.run(_run())
 
@@ -113,8 +140,8 @@ async def _run():
         store.search.side_effect = Exception("search failed")
         messages = [_make_chat_message("What is the networking plan for AKS?")]
 
-        context = await provider.invoking(messages)
-        assert context.instructions is None
+        instructions = await _call_before_run(provider, messages)
+        assert instructions is None
 
     asyncio.run(_run())
 
@@ -125,7 +152,7 @@ async def _run():
         long_text = "x" * 5000
         messages = [_make_chat_message(long_text)]
 
-        await provider.invoking(messages)
+        await _call_before_run(provider, messages)
 
         query = store.search.call_args.kwargs["query"]
         assert len(query) <= 2000
@@ -142,7 +169,7 @@ async def _run():
             _make_chat_message("Latest question about storage"),
         ]
 
-        await provider.invoking(messages)
+        await _call_before_run(provider, messages)
 
         query = store.search.call_args.kwargs["query"]
         assert "Latest question about storage" in query
@@ -159,10 +186,10 @@ async def _run():
         store.search.return_value = large_memories
         messages = [_make_chat_message("What storage configuration should we use for persistent volumes?")]
 
-        context = await provider.invoking(messages)
+        instructions = await _call_before_run(provider, messages)
 
-        assert context.instructions is not None
-        assert len(context.instructions) <= MAX_MEMORY_CONTEXT_CHARS + 200
+        assert instructions is not None
+        assert len(instructions) <= MAX_MEMORY_CONTEXT_CHARS + 200
 
     asyncio.run(_run())
 
@@ -175,10 +202,10 @@ async def _run():
         ]
         messages = [_make_chat_message("What storage class should we choose for the cluster?")]
 
-        context = await provider.invoking(messages)
+        instructions = await _call_before_run(provider, messages)
 
-        assert "Chief Architect" in context.instructions
-        assert "design" in context.instructions
+        assert "Chief Architect" in instructions
+        assert "design" in instructions
 
     asyncio.run(_run())
 
@@ -189,26 +216,25 @@ async def _run():
         store.search.return_value = [_make_memory_entry("some memory")]
         single = _make_chat_message("What about networking configuration for AKS?")
 
-        context = await provider.invoking(single)
+        instructions = await _call_before_run(provider, [single])
 
-        assert context.instructions is not None
+        assert instructions is not None
         store.search.assert_called_once()
 
     asyncio.run(_run())
 
 
 # ---------------------------------------------------------------------------
-# invoked() — Post-LLM memory storage
+# after_run() — Post-LLM memory storage
 # ---------------------------------------------------------------------------
 
 
 def test_invoked_stores_response():
     async def _run():
         provider, store = _make_provider()
-        request = [_make_chat_message("What is the networking plan for AKS?")]
-        response = [_make_chat_message("We should use Azure CNI for networking configuration in the AKS cluster")]
+        response_text = "We should use Azure CNI for networking configuration in the AKS cluster"
 
-        await provider.invoked(request, response)
+        await _call_after_run(provider, response_text)
         await provider.flush()
 
         store.add.assert_called_once()
@@ -222,10 +248,9 @@ async def _run():
 def test_invoked_skips_on_exception():
     async def _run():
         provider, store = _make_provider()
-        request = [_make_chat_message("Q")]
-        response = [_make_chat_message("A" * 100)]
-
-        await provider.invoked(request, response, invoke_exception=Exception("fail"))
+        # after_run with no response simulates exception path
+        ctx = _make_context(response_text=None)
+        await provider.after_run(agent=MagicMock(), session=MagicMock(), context=ctx, state={})
         store.add.assert_not_called()
 
     asyncio.run(_run())
@@ -234,9 +259,8 @@ async def _run():
 def test_invoked_skips_none_response():
     async def _run():
         provider, store = _make_provider()
-        request = [_make_chat_message("Q")]
-
-        await provider.invoked(request, None)
+        ctx = _make_context(response_text=None)
+        await provider.after_run(agent=MagicMock(), session=MagicMock(), context=ctx, state={})
         store.add.assert_not_called()
 
     asyncio.run(_run())
@@ -245,10 +269,8 @@ async def _run():
 def test_invoked_skips_short_response():
     async def _run():
         provider, store = _make_provider()
-        request = [_make_chat_message("Q")]
-        short = [_make_chat_message("x" * (MIN_CONTENT_LENGTH_TO_STORE - 1))]
-
-        await provider.invoked(request, short)
+        short_text = "x" * (MIN_CONTENT_LENGTH_TO_STORE - 1)
+        await _call_after_run(provider, short_text)
         store.add.assert_not_called()
 
     asyncio.run(_run())
@@ -257,10 +279,8 @@ async def _run():
 def test_invoked_stores_long_response():
     async def _run():
         provider, store = _make_provider()
-        request = [_make_chat_message("Q")]
-        long_resp = [_make_chat_message("x" * (MIN_CONTENT_LENGTH_TO_STORE + 1))]
-
-        await provider.invoked(request, long_resp)
+        long_text = "x" * (MIN_CONTENT_LENGTH_TO_STORE + 1)
+        await _call_after_run(provider, long_text)
         await provider.flush()
         store.add.assert_called_once()
 
@@ -270,11 +290,10 @@ async def _run():
 def test_invoked_increments_turn_counter():
     async def _run():
         provider, store = _make_provider()
-        request = [_make_chat_message("Q")]
-        response = [_make_chat_message("A" * 100)]
+        response_text = "A" * 100
 
-        await provider.invoked(request, response)
-        await provider.invoked(request, response)
+        await _call_after_run(provider, response_text)
+        await _call_after_run(provider, response_text)
         assert provider._turn_counter == 2
 
     asyncio.run(_run())
@@ -284,10 +303,9 @@ def test_invoked_store_failure_does_not_raise():
     async def _run():
         provider, store = _make_provider()
         store.add.side_effect = Exception("store failed")
-        request = [_make_chat_message("Q")]
-        response = [_make_chat_message("A" * 100)]
+        response_text = "A" * 100
 
-        await provider.invoked(request, response)
+        await _call_after_run(provider, response_text)
         await provider.flush()  # Should not raise
 
     asyncio.run(_run())
@@ -296,10 +314,9 @@ async def _run():
 def test_invoked_with_single_message():
     async def _run():
         provider, store = _make_provider()
-        request = _make_chat_message("What is the question about networking?")
-        response = _make_chat_message("We should use Azure CNI Overlay for the networking configuration in AKS")
+        response_text = "We should use Azure CNI Overlay for the networking configuration in AKS"
 
-        await provider.invoked(request, response)
+        await _call_after_run(provider, response_text)
         await provider.flush()
         store.add.assert_called_once()
 

src/processor/src/tests/unit/steps/test_migration_processor_run.py

@@ -54,12 +54,12 @@ def _make_processor(events: list, memory_store=None) -> MigrationProcessor:
 
     proc._telemetry = telemetry  # expose for assertions
 
-    async def _stream(_input):
+    async def _stream(_input, **kwargs):
         for ev in events:
             yield ev
 
     workflow = MagicMock()
-    workflow.run_stream = _stream
+    workflow.run = _stream
     proc.workflow = workflow
 
     # Patch _create_memory_store as an AsyncMock returning the provided value.