Spaces:
Paused
Paused
| import os | |
| import sys | |
| import unittest | |
| from unittest.mock import AsyncMock, MagicMock, patch | |
| # Add repository root to sys.path | |
| sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) | |
| from browser_utils import auth_rotation | |
| from config.global_state import GlobalState | |
| class TestAuthRotationCooldownFix(unittest.IsolatedAsyncioTestCase): | |
| async def asyncSetUp(self): | |
| # Reset GlobalState | |
| GlobalState.reset_quota_status() | |
| GlobalState.current_profile_exhausted_models = set() | |
| GlobalState.last_error_type = None | |
| GlobalState.AUTH_ROTATION_LOCK.set() # Ensure lock is open initially | |
| # Create mock state object (replacing old server mock) | |
| self.mock_state = MagicMock() | |
| self.mock_state.current_auth_profile_path = ( | |
| "auth_profiles/saved/old_profile.json" | |
| ) | |
| self.mock_state.current_ai_studio_model_id = "gemini-1.5-pro" | |
| self.mock_state.page_instance = MagicMock() | |
| self.mock_state.page_instance.is_closed.return_value = False | |
| self.mock_state.page_instance.context = MagicMock() | |
| self.mock_state.page_instance.context.clear_cookies = AsyncMock() | |
| self.mock_state.page_instance.context.add_cookies = AsyncMock() | |
| # Apply state mock (auth_rotation now uses 'state' from api_utils.server_state) | |
| self.state_patcher = patch("browser_utils.auth_rotation.state", self.mock_state) | |
| self.state_patcher.start() | |
| # Mock file operations to avoid actual IO | |
| self.open_patcher = patch("builtins.open", new_callable=MagicMock) | |
| self.mock_open = self.open_patcher.start() | |
| self.json_load_patcher = patch("json.load", return_value={"cookies": []}) | |
| self.json_load_patcher.start() | |
| # Mock os.path.exists to always return True | |
| self.exists_patcher = patch("os.path.exists", return_value=True) | |
| self.exists_patcher.start() | |
| async def asyncTearDown(self): | |
| self.state_patcher.stop() | |
| self.open_patcher.stop() | |
| self.json_load_patcher.stop() | |
| self.exists_patcher.stop() | |
| async def test_model_specific_cooldown_application( | |
| self, mock_canary, mock_get_next, mock_save | |
| ): | |
| """Test that quota exhaustion results in model-specific cooldowns, not just 'default'.""" | |
| # Setup scenario: Quota Exceeded for a specific model | |
| GlobalState.last_error_type = "QUOTA_EXCEEDED" | |
| GlobalState.current_profile_exhausted_models.add("gemini-1.5-pro") | |
| # Mock the cooldown profiles dict in the module | |
| # We need to make sure it's a real dict so updates persist during the function call | |
| test_cooldown_profiles = {} | |
| with patch( | |
| "browser_utils.auth_rotation._COOLDOWN_PROFILES", test_cooldown_profiles | |
| ): | |
| # Perform rotation | |
| await auth_rotation.perform_auth_rotation(target_model_id="gemini-1.5-pro") | |
| # Verify save was called | |
| mock_save.assert_called() | |
| # Verify the correct structure was saved | |
| # Expected: {'auth_profiles/saved/old_profile.json': {'gemini-1.5-pro': timestamp}} | |
| saved_data = mock_save.call_args[0][0] | |
| old_profile = "auth_profiles/saved/old_profile.json" | |
| self.assertIn(old_profile, saved_data) | |
| self.assertIsInstance(saved_data[old_profile], dict) | |
| self.assertIn("gemini-1.5-pro", saved_data[old_profile]) | |
| # Ensure 'default' or 'global' is NOT present if only specific model failed | |
| self.assertNotIn("default", saved_data[old_profile]) | |
| self.assertNotIn("global", saved_data[old_profile]) | |
| async def test_fallback_to_current_model_id( | |
| self, mock_canary, mock_get_next, mock_save | |
| ): | |
| """Test fallback to server.current_ai_studio_model_id if exhausted set is empty.""" | |
| # Setup scenario: Quota Exceeded, but set is empty (e.g. detected via UI text, not token count) | |
| GlobalState.last_error_type = "QUOTA_EXCEEDED" | |
| GlobalState.current_profile_exhausted_models = set() | |
| # Ensure state has a current model ID | |
| self.mock_state.current_ai_studio_model_id = "gemini-ultra" | |
| test_cooldown_profiles = {} | |
| with patch( | |
| "browser_utils.auth_rotation._COOLDOWN_PROFILES", test_cooldown_profiles | |
| ): | |
| await auth_rotation.perform_auth_rotation() | |
| saved_data = mock_save.call_args[0][0] | |
| old_profile = "auth_profiles/saved/old_profile.json" | |
| self.assertIn(old_profile, saved_data) | |
| self.assertIsInstance(saved_data[old_profile], dict) | |
| # Should have used the fallback model ID | |
| self.assertIn("gemini-ultra", saved_data[old_profile]) | |
| if __name__ == "__main__": | |
| unittest.main() | |