test(pre_init): replace assert len(mock_req.request_history) == N checks with specific request filtering

Signed-off-by: Teo <[email protected]>

Author: teocns

tests/unit/test_pre_init.py

@@ -29,16 +29,23 @@ def test_track_agent(self, mock_req):
         agentops.init(api_key=self.api_key)
         time.sleep(1)
 
-        # Assert
-        # start session and create agent
-        agentops.end_session(end_state="Success")
+        # Find agent creation request
+        agent_requests = [r for r in mock_req.request_history if "/v2/create_agent" in r.url]
+        assert len(agent_requests) > 0
+        last_agent_request = agent_requests[-1]
+
+        # Assert agent creation
+        assert last_agent_request.headers["X-Agentops-Api-Key"] == self.api_key
 
-        # Wait for flush
+        # End session and wait for flush
+        agentops.end_session(end_state="Success")
         time.sleep(1.5)
 
-        # 4 requests: check_for_updates, create_session, create_agent, update_session
-        assert len(mock_req.request_history) == 4
+        # Find session end request
+        end_session_requests = [r for r in mock_req.request_history if "/v2/update_session" in r.url]
+        assert len(end_session_requests) > 0
+        last_end_request = end_session_requests[-1]
 
-        assert mock_req.request_history[-2].headers["X-Agentops-Api-Key"] == self.api_key
+        assert last_end_request.headers["X-Agentops-Api-Key"] == self.api_key
 
         mock_req.reset()