diff --git a/tests/test_chatgpt_telegram_inference_bot.py b/tests/test_chatgpt_telegram_inference_bot.py new file mode 100644 index 0000000..a0f0bdb --- /dev/null +++ b/tests/test_chatgpt_telegram_inference_bot.py @@ -0,0 +1,158 @@ +import unittest +from unittest.mock import MagicMock, patch, ANY +import os + +# Assuming chatgpt_telegram_inference_bot.py and its parent are accessible +from chatgpt_telegram_inference_bot import ChatGPTTelegramInferenceBot +from openai_compatible_inference_bot import OpenAICompatibleInferenceBot # For patching super + +class TestChatGPTTelegramInferenceBot(unittest.IsolatedAsyncioTestCase): + + def setUp(self): + # Store and clear relevant environment variables + self.original_openai_key = os.environ.get("OPENAI_API_KEY") + self.original_small_model = os.environ.get("OPENAI_SMALL_MODEL") + self.original_large_model = os.environ.get("OPENAI_LARGE_MODEL") + self.original_small_tokens = os.environ.get("OPENAI_SMALL_MODEL_MAX_TOKENS") + self.original_large_tokens = os.environ.get("OPENAI_LARGE_MODEL_MAX_TOKENS") + self.original_system_prompt_path = os.environ.get("SYSTEM_PROMPT_PATH") + + for key in ["OPENAI_API_KEY", "OPENAI_SMALL_MODEL", "OPENAI_LARGE_MODEL", + "OPENAI_SMALL_MODEL_MAX_TOKENS", "OPENAI_LARGE_MODEL_MAX_TOKENS", "SYSTEM_PROMPT_PATH"]: + if os.environ.get(key): + del os.environ[key] + + # Mock the OpenAI client that OpenAICompatibleInferenceBot's __init__ might create + self.mock_openai_client = MagicMock() + + def tearDown(self): + # Restore environment variables + if self.original_openai_key: os.environ["OPENAI_API_KEY"] = self.original_openai_key + if self.original_small_model: os.environ["OPENAI_SMALL_MODEL"] = self.original_small_model + if self.original_large_model: os.environ["OPENAI_LARGE_MODEL"] = self.original_large_model + if self.original_small_tokens: os.environ["OPENAI_SMALL_MODEL_MAX_TOKENS"] = self.original_small_tokens + if self.original_large_tokens: os.environ["OPENAI_LARGE_MODEL_MAX_TOKENS"] = self.original_large_tokens + if self.original_system_prompt_path: os.environ["SYSTEM_PROMPT_PATH"] = self.original_system_prompt_path + + + @patch.object(OpenAICompatibleInferenceBot, '__init__') # Mock the superclass's __init__ + def test_init_defaults_and_super_call(self, mock_super_init): + os.environ["OPENAI_API_KEY"] = "test_key_chatgpt" + os.environ["OPENAI_SMALL_MODEL"] = "gpt-3.5-turbo-env" + os.environ["OPENAI_SMALL_MODEL_MAX_TOKENS"] = "350" + + bot = ChatGPTTelegramInferenceBot() + + mock_super_init.assert_called_once_with( + client=None, # ChatGPT bot will let superclass create it + api_key="test_key_chatgpt", # Passed to super + base_url=None, + api_version=None, + azure_deployment=None, + model_name="gpt-3.5-turbo-env", # Default small model from env + max_tokens_str="350", # Default small model tokens from env + small_model_name="gpt-3.5-turbo-env", + small_model_max_tokens_str="350", + large_model_name=os.environ.get("OPENAI_LARGE_MODEL", "gpt-4-turbo-preview"), # Default large + large_model_max_tokens_str=os.environ.get("OPENAI_LARGE_MODEL_MAX_TOKENS"), + system_prompt_content=None, + system_prompt_path=None, + is_gemini=False, + max_history_length=20 # Default from OpenAICompatibleInferenceBot + ) + + @patch.object(OpenAICompatibleInferenceBot, '__init__') + def test_init_with_arguments(self, mock_super_init): + mock_client_arg = MagicMock() + bot = ChatGPTTelegramInferenceBot( + openai_client=mock_client_arg, + api_key="arg_key", + small_model_name="arg_small_model", + small_model_max_tokens="123", + large_model_name="arg_large_model", + large_model_max_tokens="456", + system_prompt_content="Arg prompt" + ) + mock_super_init.assert_called_once_with( + client=mock_client_arg, + api_key="arg_key", + base_url=None, + api_version=None, + azure_deployment=None, + model_name="arg_small_model", # Initially configured with small model + max_tokens_str="123", + small_model_name="arg_small_model", + small_model_max_tokens_str="123", + large_model_name="arg_large_model", + large_model_max_tokens_str="456", + system_prompt_content="Arg prompt", + system_prompt_path=None, + is_gemini=False, + max_history_length=20 + ) + + # Test switch_model - this method is part of ChatGPTTelegramInferenceBot + # It calls _configure_model_and_tokens which is in the superclass. + # We need a bot instance where _configure_model_and_tokens can be called. + @patch('openai.OpenAI') # To allow instantiation of the bot by mocking client creation + async def test_switch_model_logic(self, mock_openai_constructor): + mock_openai_constructor.return_value = self.mock_openai_client # Mock client creation in super + + # Set env vars for model names that switch_model will use as fallback + os.environ["OPENAI_SMALL_MODEL"] = "env-small-gpt" + os.environ["OPENAI_SMALL_MODEL_MAX_TOKENS"] = "100" + os.environ["OPENAI_LARGE_MODEL"] = "env-large-gpt" + os.environ["OPENAI_LARGE_MODEL_MAX_TOKENS"] = "200" + + # Instantiate with initial model (small) + bot = ChatGPTTelegramInferenceBot() + self.assertEqual(bot.model, "env-small-gpt") + self.assertEqual(bot.max_tokens, 100) + + # Switch to large + status = await bot.switch_model() + self.assertEqual(bot.model, "env-large-gpt") + self.assertEqual(bot.max_tokens, 200) + self.assertEqual(status, "Switched to model: env-large-gpt") + + # Switch back to small + status = await bot.switch_model() + self.assertEqual(bot.model, "env-small-gpt") + self.assertEqual(bot.max_tokens, 100) + self.assertEqual(status, "Switched to model: env-small-gpt") + + @patch('openai.OpenAI') + async def test_switch_model_uses_instance_configs_if_provided(self, mock_openai_constructor): + mock_openai_constructor.return_value = self.mock_openai_client + + # Instantiate with specific model names, overriding potential env vars + bot = ChatGPTTelegramInferenceBot( + small_model_name="init-small", small_model_max_tokens="50", + large_model_name="init-large", large_model_max_tokens="150" + ) + self.assertEqual(bot.model, "init-small") # Starts with small + self.assertEqual(bot.max_tokens, 50) + + # Switch to large + status = await bot.switch_model() + self.assertEqual(bot.model, "init-large") + self.assertEqual(bot.max_tokens, 150) + self.assertEqual(status, "Switched to model: init-large") + + # Switch back to small + status = await bot.switch_model() + self.assertEqual(bot.model, "init-small") + self.assertEqual(bot.max_tokens, 50) + self.assertEqual(status, "Switched to model: init-small") + + # get_llm_description is inherited from OpenAICompatibleInferenceBot. + # Test just to ensure it works in the context of a ChatGPTBot instance + @patch('openai.OpenAI') + def test_get_llm_description_for_chatgpt_bot(self, mock_openai_constructor): + mock_openai_constructor.return_value = self.mock_openai_client + bot = ChatGPTTelegramInferenceBot(small_model_name="gpt-3.5-desc", small_model_max_tokens="777") + # Initially configured with small model + self.assertEqual(bot.get_llm_description(), "LLM: gpt-3.5-desc, Max Tokens: 777, Azure: False") + +if __name__ == '__main__': + unittest.main()