emma-hermes/tests/cron/test_cron_usage_attribution.py

156 lines
4.7 KiB
Python

import sys
from pathlib import Path
from unittest.mock import MagicMock, patch
# Ensure project root is importable.
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
import cron.scheduler as scheduler
def _base_job(**overrides):
job = {
"id": "cron-usage-test",
"name": "cron usage test",
"prompt": "hello",
"model": None,
"provider": None,
"provider_snapshot": None,
"model_snapshot": None,
"base_url": None,
"deliver": "local",
}
job.update(overrides)
return job
def test_cron_usage_source_infers_telegram_dm_user_id_from_chat_id():
source = scheduler._cron_usage_source_from_job(
{
"origin": {
"platform": "telegram",
"chat_id": "601504103",
"chat_name": "tester",
}
}
)
assert source is not None
assert source.platform.value == "telegram"
assert source.chat_id == "601504103"
assert source.user_id == "601504103"
def test_cron_usage_source_does_not_infer_telegram_group_chat_id():
source = scheduler._cron_usage_source_from_job(
{
"origin": {
"platform": "telegram",
"chat_id": "-1001234567890",
"chat_name": "group",
}
}
)
assert source is None
def test_cron_usage_source_infers_line_user_id_from_dm_chat_id():
source = scheduler._cron_usage_source_from_job(
{
"origin": {
"platform": "line",
"chat_id": "U1234567890abcdef",
"chat_name": "line user",
}
}
)
assert source is not None
assert source.platform.value == "line"
assert source.chat_id == "U1234567890abcdef"
assert source.user_id == "U1234567890abcdef"
def test_run_job_records_usage_for_origin_principal(tmp_path, monkeypatch):
job = _base_job(
origin={
"platform": "telegram",
"chat_id": "601504103",
"user_id": "601504103",
"chat_name": "tester",
}
)
fake_db = MagicMock()
principal_usage_calls = []
runtime_usage_calls = []
class FakeGatewayUserStore:
def record_usage_for_source(self, source, **kwargs):
principal_usage_calls.append((source, kwargs))
return {"principal_id": "p-test"}
def fake_runtime_usage(source, **kwargs):
runtime_usage_calls.append((source, kwargs))
return {"allowed": True}
with patch("cron.scheduler._hermes_home", tmp_path), \
patch("hermes_cli.env_loader.load_hermes_dotenv"), \
patch("hermes_cli.env_loader.reset_secret_source_cache"), \
patch("hermes_state.SessionDB", return_value=fake_db), \
patch(
"hermes_cli.runtime_provider.resolve_runtime_provider",
return_value={
"api_key": "test-key",
"base_url": "https://example.invalid/v1",
"provider": "openrouter",
"api_mode": "chat_completions",
},
), \
patch("run_agent.AIAgent") as mock_agent_cls, \
patch("gateway.user_verification.GatewayUserStore", FakeGatewayUserStore), \
patch("gateway.runtime_governance.record_runtime_usage_for_source", fake_runtime_usage):
mock_agent = MagicMock()
mock_agent.run_conversation.return_value = {
"final_response": "ok",
"completed": True,
"failed": False,
"turn_exit_reason": "assistant_response",
}
mock_agent.session_input_tokens = 111
mock_agent.session_output_tokens = 22
mock_agent.session_total_tokens = 133
mock_agent.model = "gpt-5.4"
mock_agent.session_id = "cron-session-1"
mock_agent_cls.return_value = mock_agent
success, output, final_response, error = scheduler.run_job(job)
assert success is True
assert error is None
assert final_response == "ok"
assert len(principal_usage_calls) == 1
source, usage_kwargs = principal_usage_calls[0]
assert source.platform.value == "telegram"
assert source.user_id == "601504103"
assert usage_kwargs == {
"input_tokens": 111,
"output_tokens": 22,
"total_tokens": 133,
"message_count": 1,
"model": "gpt-5.4",
}
assert len(runtime_usage_calls) == 1
runtime_source, runtime_kwargs = runtime_usage_calls[0]
assert runtime_source.platform.value == "telegram"
assert runtime_source.user_id == "601504103"
assert runtime_kwargs == {
"input_tokens": 111,
"output_tokens": 22,
"model": "gpt-5.4",
"session_id": "cron-session-1",
"question_count": 1,
}