test_compressor_fallback_update.py
1 """Tests that _try_activate_fallback updates the context compressor.""" 2 3 from unittest.mock import MagicMock, patch 4 5 from run_agent import AIAgent 6 from agent.context_compressor import ContextCompressor 7 8 9 def _make_agent_with_compressor() -> AIAgent: 10 """Build a minimal AIAgent with a context_compressor, skipping __init__.""" 11 agent = AIAgent.__new__(AIAgent) 12 13 # Primary model settings 14 agent.model = "primary-model" 15 agent.provider = "openrouter" 16 agent.base_url = "https://openrouter.ai/api/v1" 17 agent.api_key = "sk-primary" 18 agent.api_mode = "chat_completions" 19 agent.client = MagicMock() 20 agent.quiet_mode = True 21 22 # Fallback config 23 agent._fallback_activated = False 24 agent._fallback_model = { 25 "provider": "openai", 26 "model": "gpt-4o", 27 } 28 agent._fallback_chain = [agent._fallback_model] 29 agent._fallback_index = 0 30 31 # Context compressor with primary model values 32 compressor = ContextCompressor( 33 model="primary-model", 34 threshold_percent=0.50, 35 base_url="https://openrouter.ai/api/v1", 36 api_key="sk-primary", 37 provider="openrouter", 38 quiet_mode=True, 39 ) 40 agent.context_compressor = compressor 41 42 return agent 43 44 45 @patch("agent.auxiliary_client.resolve_provider_client") 46 @patch("agent.model_metadata.get_model_context_length", return_value=128_000) 47 def test_compressor_updated_on_fallback(mock_ctx_len, mock_resolve): 48 """After fallback activation, the compressor must reflect the fallback model.""" 49 agent = _make_agent_with_compressor() 50 51 assert agent.context_compressor.model == "primary-model" 52 53 fb_client = MagicMock() 54 fb_client.base_url = "https://api.openai.com/v1" 55 fb_client.api_key = "sk-fallback" 56 mock_resolve.return_value = (fb_client, None) 57 58 agent._is_direct_openai_url = lambda url: "api.openai.com" in url 59 agent._emit_status = lambda msg: None 60 61 result = agent._try_activate_fallback() 62 63 assert result is True 64 assert agent._fallback_activated is True 65 66 c = agent.context_compressor 67 assert c.model == "gpt-4o" 68 assert c.base_url == "https://api.openai.com/v1" 69 assert c.api_key == "sk-fallback" 70 assert c.provider == "openai" 71 assert c.context_length == 128_000 72 assert c.threshold_tokens == int(128_000 * c.threshold_percent) 73 74 75 @patch("agent.auxiliary_client.resolve_provider_client") 76 @patch("agent.model_metadata.get_model_context_length", return_value=128_000) 77 def test_compressor_not_present_does_not_crash(mock_ctx_len, mock_resolve): 78 """If the agent has no compressor, fallback should still succeed.""" 79 agent = _make_agent_with_compressor() 80 agent.context_compressor = None 81 82 fb_client = MagicMock() 83 fb_client.base_url = "https://api.openai.com/v1" 84 fb_client.api_key = "sk-fallback" 85 mock_resolve.return_value = (fb_client, None) 86 87 agent._is_direct_openai_url = lambda url: "api.openai.com" in url 88 agent._emit_status = lambda msg: None 89 90 result = agent._try_activate_fallback() 91 assert result is True