/ tests / run_agent / test_compressor_fallback_update.py
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