Squash merge teocns/improve-session-tests-assertions into logcap

Signed-off-by: Teo <[email protected]>
AgentOps-AI · Feb 6, 2025 · ec1978c · ec1978c
1 parent 76551b6
commit ec1978c
Showing 1 changed file with 24 additions and 22 deletions.
diff --git a/tests/unit/test_record_tool.py b/tests/unit/test_record_tool.py
@@ -108,10 +108,14 @@ async def async_add(x, y):
 
         # Assert
         assert result == 7
-        # Assert
-        assert len(mock_req.request_history) == 3
-        assert mock_req.last_request.headers["X-Agentops-Api-Key"] == self.api_key
-        request_json = mock_req.last_request.json()
+
+        # Find the record_tool request
+        tool_requests = [r for r in mock_req.request_history if "/v2/create_events" in r.url]
+        assert len(tool_requests) > 0
+        last_tool_request = tool_requests[-1]
+
+        assert last_tool_request.headers["X-Agentops-Api-Key"] == self.api_key
+        request_json = last_tool_request.json()
         assert request_json["events"][0]["name"] == self.tool_name
         assert request_json["events"][0]["params"] == {"x": 3, "y": 4}
         assert request_json["events"][0]["returns"] == 7
@@ -186,31 +190,29 @@ async def async_add(x, y):
         await async_add(1, 2, session=session_2)
         time.sleep(0.1)
 
-        # Assert
-        assert len(mock_req.request_history) == 5
+        # Find tool requests
+        tool_requests = [r for r in mock_req.request_history if "/v2/create_events" in r.url]
+        assert len(tool_requests) >= 2  # Should have at least 2 tool requests
 
-        request_json = mock_req.last_request.json()
-        assert mock_req.last_request.headers["X-Agentops-Api-Key"] == self.api_key
-        assert (
-            mock_req.last_request.headers["Authorization"]
-            == f"Bearer {mock_req.session_jwts[str(session_2.session_id)]}"
-        )
+        # Verify session_2's request (last request)
+        last_request = tool_requests[-1]
+        assert last_request.headers["X-Agentops-Api-Key"] == self.api_key
+        assert last_request.headers["Authorization"] == f"Bearer {mock_req.session_jwts[str(session_2.session_id)]}"
+        request_json = last_request.json()
         assert request_json["events"][0]["name"] == self.tool_name
         assert request_json["events"][0]["params"] == {"x": 1, "y": 2}
         assert request_json["events"][0]["returns"] == 3
 
-        second_last_request_json = mock_req.request_history[-2].json()
-        assert mock_req.request_history[-2].headers["X-Agentops-Api-Key"] == self.api_key
+        # Verify session_1's request (second to last request)
+        second_last_request = tool_requests[-2]
+        assert second_last_request.headers["X-Agentops-Api-Key"] == self.api_key
         assert (
-            mock_req.request_history[-2].headers["Authorization"]
-            == f"Bearer {mock_req.session_jwts[str(session_1.session_id)]}"
+            second_last_request.headers["Authorization"] == f"Bearer {mock_req.session_jwts[str(session_1.session_id)]}"
         )
-        assert second_last_request_json["events"][0]["name"] == self.tool_name
-        assert second_last_request_json["events"][0]["params"] == {
-            "x": 1,
-            "y": 2,
-        }
-        assert second_last_request_json["events"][0]["returns"] == 3
+        request_json = second_last_request.json()
+        assert request_json["events"][0]["name"] == self.tool_name
+        assert request_json["events"][0]["params"] == {"x": 1, "y": 2}
+        assert request_json["events"][0]["returns"] == 3
 
         session_1.end_session(end_state="Success")
         session_2.end_session(end_state="Success")