From 2c667b4cb0e2b0c4e9ae10ee7544ac9b7cbfb53d Mon Sep 17 00:00:00 2001
From: Jeffrey Phillips Freeman <the@jeffreyfreeman.me>
Date: Mon, 30 Jun 2025 23:16:09 +0000
Subject: [PATCH] test: enhance message routing verification with context
 history

---
 .../steps/router_comprehensive_steps.py       | 32 ++++++++++++++++---
 1 file changed, 27 insertions(+), 5 deletions(-)

diff --git a/tests/features/steps/router_comprehensive_steps.py b/tests/features/steps/router_comprehensive_steps.py
index cb3e4d889..103ad60e0 100644
--- a/tests/features/steps/router_comprehensive_steps.py
+++ b/tests/features/steps/router_comprehensive_steps.py
@@ -87,31 +87,53 @@ def step_impl(context):
 
 @then("each message should be routed to the appropriate agent")
 def step_impl(context):
-    # Check that each message was routed to the right agent
+    # Verify that each message was routed to the correct agent
     assert context.responses["help"] == "Response from agent1"
     assert context.responses["search"] == "Response from agent2"
     assert context.responses["long"] == "Response from agent3"
 
-    # Check that each agent's process method was called with the right message and context
+    # Agent1's context history will contain its own interaction
     context.agents["agent1"].process.assert_called_once_with(
         "I need help with something",
         {
-            "message": "I need help with something",
             "initial_message": "I need help with something",
+            "history": [
+                {
+                    "source": "input",
+                    "destination": "agent1",
+                    "message": "Response from agent1",
+                }
+            ],
         },
     )
+
+    # Agent2's context history will contain its own interaction
     context.agents["agent2"].process.assert_called_once_with(
         "Can you search for information?",
         {
-            "message": "Can you search for information?",
             "initial_message": "Can you search for information?",
+            "history": [
+                {
+                    "source": "input",
+                    "destination": "agent2",
+                    "message": "Response from agent2",
+                }
+            ],
         },
     )
+
+    # Agent3's context history will contain its own interaction
     context.agents["agent3"].process.assert_called_once_with(
         "This is a longer message with more than five words",
         {
-            "message": "This is a longer message with more than five words",
             "initial_message": "This is a longer message with more than five words",
+            "history": [
+                {
+                    "source": "input",
+                    "destination": "agent3",
+                    "message": "Response from agent3",
+                },
+            ],
         },
     )