1  from types import SimpleNamespace
 2  from unittest.mock import MagicMock, patch
 3  
 4  from run_agent import AIAgent
 5  
 6  
 7  def _mock_response(*, usage: dict, content: str = "done"):
 8      msg = SimpleNamespace(content=content, tool_calls=None)
 9      choice = SimpleNamespace(message=msg, finish_reason="stop")
10      return SimpleNamespace(
11          choices=[choice],
12          model="test/model",
13          usage=SimpleNamespace(**usage),
14      )
15  
16  
17  def _make_agent(session_db, *, platform: str):
18      with (
19          patch("run_agent.get_tool_definitions", return_value=[]),
20          patch("run_agent.check_toolset_requirements", return_value={}),
21          patch("run_agent.OpenAI"),
22      ):
23          agent = AIAgent(
24              api_key="test-key",
25              base_url="https://openrouter.ai/api/v1",
26              quiet_mode=True,
27              skip_context_files=True,
28              skip_memory=True,
29              session_db=session_db,
30              session_id=f"{platform}-session",
31              platform=platform,
32          )
33      agent.client = MagicMock()
34      agent.client.chat.completions.create.return_value = _mock_response(
35          usage={
36              "prompt_tokens": 11,
37              "completion_tokens": 7,
38              "total_tokens": 18,
39          }
40      )
41      return agent
42  
43  
44  def test_run_conversation_persists_tokens_for_telegram_sessions():
45      session_db = MagicMock()
46      agent = _make_agent(session_db, platform="telegram")
47  
48      result = agent.run_conversation("hello")
49  
50      assert result["final_response"] == "done"
51      session_db.update_token_counts.assert_called_once()
52      assert session_db.update_token_counts.call_args.args[0] == "telegram-session"
53  
54  
55  def test_run_conversation_persists_tokens_for_cron_sessions():
56      session_db = MagicMock()
57      agent = _make_agent(session_db, platform="cron")
58  
59      result = agent.run_conversation("hello")
60  
61      assert result["final_response"] == "done"
62      session_db.update_token_counts.assert_called_once()
63      assert session_db.update_token_counts.call_args.args[0] == "cron-session"