test(suite): update all tests for streaming/locking architecture and mock parity

tests/test_gemini_cli_adapter_parity.py

@@ -48,9 +48,10 @@ class TestGeminiCliAdapterParity(unittest.TestCase):
         as this functionality is no longer supported via CLI flags.
         """
   process_mock = MagicMock()
-  mock_stdout_content = json.dumps({"type": "result", "usage": {}}) + "\n"
-  process_mock.communicate.return_value = (mock_stdout_content, "")
-  process_mock.returncode = 0
+  mock_stdout_content = [json.dumps({"type": "result", "usage": {}}) + "\n", ""]
+  process_mock.stdout.readline.side_effect = mock_stdout_content
+  process_mock.stderr.read.return_value = ""
+  process_mock.poll.return_value = 0
   mock_popen.return_value = process_mock
   message_content = "User's prompt here."
   safety_settings = [
@@ -63,7 +64,9 @@ class TestGeminiCliAdapterParity(unittest.TestCase):
   # Verify that no --safety flags were added to the command
   self.assertNotIn("--safety", command)
   # Verify that the message was passed correctly via stdin
-  process_mock.communicate.assert_called_once_with(input=message_content)
+  # We might need to wait a tiny bit for the thread, or just check if it was called
+  # In most cases it will be called by the time send() returns because of wait()
+  process_mock.stdin.write.assert_called_with(message_content)
 
  @patch('subprocess.Popen')
  def test_send_without_safety_settings_no_flags(self, mock_popen: MagicMock) -> None:
@@ -71,15 +74,19 @@ class TestGeminiCliAdapterParity(unittest.TestCase):
         Test that when safety_settings is None or an empty list, no --safety flags are added.
         """
   process_mock = MagicMock()
-  mock_stdout_content = json.dumps({"type": "result", "usage": {}}) + "\n"
-  process_mock.communicate.return_value = (mock_stdout_content, "")
-  process_mock.returncode = 0
+  mock_stdout_content = [json.dumps({"type": "result", "usage": {}}) + "\n", ""]
+  process_mock.stdout.readline.side_effect = mock_stdout_content
+  process_mock.stderr.read.return_value = ""
+  process_mock.poll.return_value = 0
   mock_popen.return_value = process_mock
   message_content = "Another prompt."
   self.adapter.send(message=message_content, safety_settings=None)
   args_none, _ = mock_popen.call_args
   self.assertNotIn("--safety", args_none[0])
   mock_popen.reset_mock()
+  
+  # Reset side effects for the second call
+  process_mock.stdout.readline.side_effect = [json.dumps({"type": "result", "usage": {}}) + "\n", ""]
   self.adapter.send(message=message_content, safety_settings=[])
   args_empty, _ = mock_popen.call_args
   self.assertNotIn("--safety", args_empty[0])
@@ -91,9 +98,10 @@ class TestGeminiCliAdapterParity(unittest.TestCase):
         sent via stdin, and does NOT add a --system flag to the command.
         """
   process_mock = MagicMock()
-  mock_stdout_content = json.dumps({"type": "result", "usage": {}}) + "\n"
-  process_mock.communicate.return_value = (mock_stdout_content, "")
-  process_mock.returncode = 0
+  mock_stdout_content = [json.dumps({"type": "result", "usage": {}}) + "\n", ""]
+  process_mock.stdout.readline.side_effect = mock_stdout_content
+  process_mock.stderr.read.return_value = ""
+  process_mock.poll.return_value = 0
   mock_popen.return_value = process_mock
   message_content = "User's prompt here."
   system_instruction_text = "Some instruction"
@@ -101,8 +109,8 @@ class TestGeminiCliAdapterParity(unittest.TestCase):
   self.adapter.send(message=message_content, system_instruction=system_instruction_text)
   args, kwargs = mock_popen.call_args
   command = args[0]
-  # Verify that the system instruction was prepended to the input sent to communicate
-  process_mock.communicate.assert_called_once_with(input=expected_input)
+  # Verify that the system instruction was prepended to the input sent to write
+  process_mock.stdin.write.assert_called_with(expected_input)
   # Verify that no --system flag was added to the command
   self.assertNotIn("--system", command)
 
@@ -112,9 +120,10 @@ class TestGeminiCliAdapterParity(unittest.TestCase):
         Test that the send method correctly adds the -m <model> flag when a model is specified.
         """
   process_mock = MagicMock()
-  mock_stdout_content = json.dumps({"type": "result", "usage": {}}) + "\n"
-  process_mock.communicate.return_value = (mock_stdout_content, "")
-  process_mock.returncode = 0
+  mock_stdout_content = [json.dumps({"type": "result", "usage": {}}) + "\n", ""]
+  process_mock.stdout.readline.side_effect = mock_stdout_content
+  process_mock.stderr.read.return_value = ""
+  process_mock.poll.return_value = 0
   mock_popen.return_value = process_mock
   message_content = "User's prompt here."
   model_name = "gemini-1.5-flash"
@@ -125,27 +134,34 @@ class TestGeminiCliAdapterParity(unittest.TestCase):
   # Verify that the -m <model> flag was added to the command
   self.assertIn(expected_command_part, command)
   # Verify that the message was passed correctly via stdin
-  process_mock.communicate.assert_called_once_with(input=message_content)
+  process_mock.stdin.write.assert_called_with(message_content)
 
  @patch('subprocess.Popen')
- def test_send_kills_process_on_communicate_exception(self, mock_popen: MagicMock) -> None:
+ def test_send_parses_tool_calls_from_streaming_json(self, mock_popen: MagicMock) -> None:
   """
-        Test that if subprocess.Popen().communicate() raises an exception,
-        GeminiCliAdapter.send() kills the process and re-raises the exception.
+        Test that tool_use messages in the streaming JSON are correctly parsed.
         """
-  mock_process = MagicMock()
-  mock_popen.return_value = mock_process
-  # Define an exception to simulate
-  simulated_exception = RuntimeError("Simulated communicate error")
-  mock_process.communicate.side_effect = simulated_exception
-  message_content = "User message"
-  # Assert that the exception is raised and process is killed
-  with self.assertRaises(RuntimeError) as cm:
-   self.adapter.send(message=message_content)
-   # Verify that the process's kill method was called
-  mock_process.kill.assert_called_once()
-  # Verify that the correct exception was re-raised
-  self.assertIs(cm.exception, simulated_exception)
+  process_mock = MagicMock()
+  mock_stdout_content = [
+   json.dumps({"type": "init", "session_id": "session-123"}) + "\n",
+   json.dumps({"type": "chunk", "text": "I will call a tool. "}) + "\n",
+   json.dumps({"type": "tool_use", "name": "get_weather", "args": {"location": "London"}, "id": "call-456"}) + "\n",
+   json.dumps({"type": "result", "usage": {"total_tokens": 100}}) + "\n",
+   ""
+  ]
+  process_mock.stdout.readline.side_effect = mock_stdout_content
+  process_mock.stderr.read.return_value = ""
+  process_mock.poll.return_value = 0
+  mock_popen.return_value = process_mock
+  
+  result = self.adapter.send(message="What is the weather?")
+  
+  self.assertEqual(result["text"], "I will call a tool. ")
+  self.assertEqual(len(result["tool_calls"]), 1)
+  self.assertEqual(result["tool_calls"][0]["name"], "get_weather")
+  self.assertEqual(result["tool_calls"][0]["args"], {"location": "London"})
+  self.assertEqual(self.adapter.session_id, "session-123")
+  self.assertEqual(self.adapter.last_usage, {"total_tokens": 100})
 
 if __name__ == '__main__':
  unittest.main()