mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-05-22 16:06:50 +00:00
Lawrance_YXLiao
2eeb597985
* fix(runs): expose active progress counters * fix(runs): avoid delayed progress flush on completion * fix(runs): tighten progress snapshot semantics * fix(runs): preserve omitted progress fields * chore(runs): remove duplicate journal initialization
83 lines
2.5 KiB
Python
83 lines
2.5 KiB
Python
"""Tests for thread-level token usage aggregation API."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
from _router_auth_helpers import make_authed_test_app
|
|
from fastapi.testclient import TestClient
|
|
|
|
from app.gateway.routers import thread_runs
|
|
|
|
|
|
def _make_app(run_store: MagicMock):
|
|
app = make_authed_test_app()
|
|
app.include_router(thread_runs.router)
|
|
app.state.run_store = run_store
|
|
return app
|
|
|
|
|
|
def test_thread_token_usage_returns_stable_shape():
|
|
run_store = MagicMock()
|
|
run_store.aggregate_tokens_by_thread = AsyncMock(
|
|
return_value={
|
|
"total_tokens": 150,
|
|
"total_input_tokens": 90,
|
|
"total_output_tokens": 60,
|
|
"total_runs": 2,
|
|
"by_model": {"unknown": {"tokens": 150, "runs": 2}},
|
|
"by_caller": {
|
|
"lead_agent": 120,
|
|
"subagent": 25,
|
|
"middleware": 5,
|
|
},
|
|
},
|
|
)
|
|
app = _make_app(run_store)
|
|
|
|
with TestClient(app) as client:
|
|
response = client.get("/api/threads/thread-1/token-usage")
|
|
|
|
assert response.status_code == 200
|
|
assert response.json() == {
|
|
"thread_id": "thread-1",
|
|
"total_tokens": 150,
|
|
"total_input_tokens": 90,
|
|
"total_output_tokens": 60,
|
|
"total_runs": 2,
|
|
"by_model": {"unknown": {"tokens": 150, "runs": 2}},
|
|
"by_caller": {
|
|
"lead_agent": 120,
|
|
"subagent": 25,
|
|
"middleware": 5,
|
|
},
|
|
}
|
|
run_store.aggregate_tokens_by_thread.assert_awaited_once_with("thread-1")
|
|
|
|
|
|
def test_thread_token_usage_can_include_active_runs():
|
|
run_store = MagicMock()
|
|
run_store.aggregate_tokens_by_thread = AsyncMock(
|
|
return_value={
|
|
"total_tokens": 175,
|
|
"total_input_tokens": 120,
|
|
"total_output_tokens": 55,
|
|
"total_runs": 3,
|
|
"by_model": {"unknown": {"tokens": 175, "runs": 3}},
|
|
"by_caller": {
|
|
"lead_agent": 145,
|
|
"subagent": 25,
|
|
"middleware": 5,
|
|
},
|
|
},
|
|
)
|
|
app = _make_app(run_store)
|
|
|
|
with TestClient(app) as client:
|
|
response = client.get("/api/threads/thread-1/token-usage?include_active=true")
|
|
|
|
assert response.status_code == 200
|
|
assert response.json()["total_tokens"] == 175
|
|
assert response.json()["total_runs"] == 3
|
|
run_store.aggregate_tokens_by_thread.assert_awaited_once_with("thread-1", include_active=True)
|