import pytest from unittest.mock import MagicMock, patch import ai_client class MockUsage: def __init__(self): self.prompt_token_count = 10 self.candidates_token_count = 5 self.total_token_count = 15 self.cached_content_token_count = 0 class MockPart: def __init__(self, text, function_call): self.text = text self.function_call = function_call class MockContent: def __init__(self, parts): self.parts = parts class MockCandidate: def __init__(self, parts): self.content = MockContent(parts) self.finish_reason = MagicMock() self.finish_reason.name = "STOP" def test_ai_client_event_emitter_exists(): # This should fail initially because 'events' won't exist on ai_client assert hasattr(ai_client, 'events') def test_event_emission(): callback = MagicMock() ai_client.events.on("test_event", callback) ai_client.events.emit("test_event", payload={"data": 123}) callback.assert_called_once_with(payload={"data": 123}) def test_send_emits_events(): with patch("ai_client._send_gemini") as mock_send_gemini, \ patch("ai_client._send_anthropic") as mock_send_anthropic: mock_send_gemini.return_value = "gemini response" start_callback = MagicMock() response_callback = MagicMock() ai_client.events.on("request_start", start_callback) ai_client.events.on("response_received", response_callback) ai_client.set_provider("gemini", "gemini-2.5-flash-lite") ai_client.send("context", "message") # We mocked _send_gemini so it doesn't emit events inside. # But wait, ai_client.send itself emits request_start and response_received? # Actually, ai_client.send delegates to _send_gemini. # Let's mock _gemini_client instead to let _send_gemini run and emit events. pass def test_send_emits_events_proper(): with patch("ai_client._ensure_gemini_client"), \ patch("ai_client._gemini_client") as mock_client: mock_chat = MagicMock() mock_client.chats.create.return_value = mock_chat mock_response = MagicMock() mock_response.candidates = [MockCandidate([MockPart("gemini response", None)])] mock_response.usage_metadata = MockUsage() mock_chat.send_message.return_value = mock_response start_callback = MagicMock() response_callback = MagicMock() ai_client.events.on("request_start", start_callback) ai_client.events.on("response_received", response_callback) ai_client.set_provider("gemini", "gemini-2.5-flash-lite") ai_client.send("context", "message") assert start_callback.called assert response_callback.called args, kwargs = start_callback.call_args assert kwargs['payload']['provider'] == 'gemini' def test_send_emits_tool_events(): import mcp_client with patch("ai_client._ensure_gemini_client"), \ patch("ai_client._gemini_client") as mock_client, \ patch("mcp_client.dispatch") as mock_dispatch: mock_chat = MagicMock() mock_client.chats.create.return_value = mock_chat # 1. Setup mock response with a tool call mock_fc = MagicMock() mock_fc.name = "read_file" mock_fc.args = {"path": "test.txt"} mock_response_with_tool = MagicMock() mock_response_with_tool.candidates = [MockCandidate([MockPart("tool call text", mock_fc)])] mock_response_with_tool.usage_metadata = MockUsage() # 2. Setup second mock response (final answer) mock_response_final = MagicMock() mock_response_final.candidates = [MockCandidate([MockPart("final answer", None)])] mock_response_final.usage_metadata = MockUsage() mock_chat.send_message.side_effect = [mock_response_with_tool, mock_response_final] mock_dispatch.return_value = "file content" ai_client.set_provider("gemini", "gemini-2.5-flash-lite") tool_callback = MagicMock() ai_client.events.on("tool_execution", tool_callback) ai_client.send("context", "message") # Should be called twice: once for 'started', once for 'completed' assert tool_callback.call_count == 2 # Check 'started' call args, kwargs = tool_callback.call_args_list[0] assert kwargs['payload']['status'] == 'started' assert kwargs['payload']['tool'] == 'read_file' # Check 'completed' call args, kwargs = tool_callback.call_args_list[1] assert kwargs['payload']['status'] == 'completed' assert kwargs['payload']['result'] == 'file content'