import unittest from unittest.mock import patch, MagicMock, ANY import json import subprocess import io import sys import os # Ensure the project root is in sys.path to resolve imports correctly project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) if project_root not in sys.path: sys.path.append(project_root) # Import the class to be tested from gemini_cli_adapter import GeminiCliAdapter class TestGeminiCliAdapterParity(unittest.TestCase): def setUp(self): """Set up a fresh adapter instance and reset session state for each test.""" # Patch session_logger to prevent file operations during tests self.session_logger_patcher = patch('gemini_cli_adapter.session_logger') self.mock_session_logger = self.session_logger_patcher.start() self.adapter = GeminiCliAdapter(binary_path="gemini") self.adapter.session_id = None self.adapter.last_usage = None self.adapter.last_latency = 0.0 def tearDown(self): self.session_logger_patcher.stop() @patch('subprocess.Popen') def test_count_tokens_uses_estimation(self, mock_popen): """ Test that count_tokens uses character-based estimation. """ contents_to_count = ["This is the first line.", "This is the second line."] expected_chars = len("\n".join(contents_to_count)) expected_tokens = expected_chars // 4 token_count = self.adapter.count_tokens(contents=contents_to_count) self.assertEqual(token_count, expected_tokens) # Verify that NO subprocess was started for counting mock_popen.assert_not_called() @patch('subprocess.Popen') def test_send_with_safety_settings_no_flags_added(self, mock_popen): """ Test that the send method does NOT add --safety flags when safety_settings are provided, 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_popen.return_value = process_mock message_content = "User's prompt here." safety_settings = [ {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_ONLY_HIGH"}, {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"} ] self.adapter.send(message=message_content, safety_settings=safety_settings) args, kwargs = mock_popen.call_args command = args[0] # 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) @patch('subprocess.Popen') def test_send_without_safety_settings_no_flags(self, mock_popen): """ 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_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() self.adapter.send(message=message_content, safety_settings=[]) args_empty, _ = mock_popen.call_args self.assertNotIn("--safety", args_empty[0]) @patch('subprocess.Popen') def test_send_with_system_instruction_prepended_to_stdin(self, mock_popen): """ Test that the send method prepends the system instruction to the prompt 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_popen.return_value = process_mock message_content = "User's prompt here." system_instruction_text = "Some instruction" expected_input = f"{system_instruction_text}\n\n{message_content}" 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 no --system flag was added to the command self.assertNotIn("--system", command) @patch('subprocess.Popen') def test_send_with_model_parameter(self, mock_popen): """ Test that the send method correctly adds the -m 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_popen.return_value = process_mock message_content = "User's prompt here." model_name = "gemini-1.5-flash" expected_command_part = f'-m "{model_name}"' self.adapter.send(message=message_content, model=model_name) args, kwargs = mock_popen.call_args command = args[0] # Verify that the -m 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) @patch('subprocess.Popen') def test_send_kills_process_on_communicate_exception(self, mock_popen): """ Test that if subprocess.Popen().communicate() raises an exception, GeminiCliAdapter.send() kills the process and re-raises the exception. """ 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) if __name__ == '__main__': unittest.main()