diff --git a/tests/integration/test_base.py b/tests/integration/test_base.py index 38f208e52..059cd0d42 100644 --- a/tests/integration/test_base.py +++ b/tests/integration/test_base.py @@ -20,9 +20,10 @@ async def async_setup_method(self, method): # Mock successful event recording response async def mock_post(url, data=None, *args, **kwargs): - mock_response = MagicMock() - mock_response.code = 200 + mock_response = AsyncMock() mock_response.status = 200 + mock_response.code = 200 + mock_response.body = b'{"status": "success"}' mock_response.json.return_value = {"status": "success"} # Store request details for verification @@ -35,7 +36,7 @@ async def mock_post(url, data=None, *args, **kwargs): return mock_response # Patch HttpClient.post to use our mock - self.http_client_patcher = patch.object(HttpClient, 'post', side_effect=mock_post) + self.http_client_patcher = patch.object(HttpClient, 'post', new=mock_post) self.mock_http_client = self.http_client_patcher.start() # Configure client and session @@ -47,9 +48,11 @@ async def mock_post(url, data=None, *args, **kwargs): async def teardown_method(self, method): """Clean up after test.""" if hasattr(self, "provider"): - self.provider.undo_override() + await self.provider.undo_override() if hasattr(self, 'http_client_patcher'): self.http_client_patcher.stop() + if hasattr(self, "session"): + await self.session.end_session() async def async_verify_events(self, session, expected_count=1): """Verify events were recorded.""" @@ -61,11 +64,11 @@ async def async_verify_events(self, session, expected_count=1): assert len(create_events_requests) >= expected_count, f"Expected at least {expected_count} event(s), but no events were recorded" return create_events_requests - async def async_verify_llm_event(self, mock_req, model=None): + async def async_verify_llm_event(self, request_history, model=None): """Verify LLM event was recorded.""" await asyncio.sleep(0.1) # Allow time for async event processing create_events_requests = [ - req for req in self.request_history + req for req in request_history if isinstance(req['url'], str) and req['url'].endswith("/v2/create_events") ] assert len(create_events_requests) >= 1, "No events were recorded"