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()