73 lines
2.7 KiB
Python
73 lines
2.7 KiB
Python
"""任务成本预估单元测试（mock 数据，不依赖数据库）。"""
|
||
from types import SimpleNamespace
|
||
from unittest.mock import patch
|
||
|
||
from app.cost_estimate import estimate_campaign_cost
|
||
|
||
|
||
class _PlannedSession:
|
||
"""按顺序返回 exec 查询结果：segments → usage_events → llm_calls。"""
|
||
|
||
def __init__(self, queues):
|
||
self.queues = list(queues)
|
||
|
||
def exec(self, *_a, **_k):
|
||
return SimpleNamespace(all=lambda: self.queues.pop(0) if self.queues else [])
|
||
|
||
|
||
def _campaign(segments=3):
|
||
return SimpleNamespace(id=1, organization_id=1, owner_user_id=7)
|
||
|
||
|
||
def _event(in_toks, out_toks, in_rate=2.0, out_rate=8.0):
|
||
return SimpleNamespace(status="measured", input_quantity=in_toks, output_quantity=out_toks,
|
||
input_rate=in_rate, output_rate=out_rate)
|
||
|
||
|
||
def test_no_segments_not_ready():
|
||
session = _PlannedSession([[]])
|
||
report = estimate_campaign_cost(session, _campaign())
|
||
assert report["ready"] is False
|
||
|
||
|
||
def test_fallback_estimate_without_history():
|
||
session = _PlannedSession([[_segment_stub() for _ in range(3)], [], []])
|
||
report = estimate_campaign_cost(session, _campaign())
|
||
assert report["ready"] is True
|
||
assert report["source"] == "fallback"
|
||
# 3 段 × (6000+2500) × 1.6 = 40800
|
||
assert report["estimated_tokens"] == 40800
|
||
assert report["estimated_points"] == 408000
|
||
assert report["estimated_amount_minor"] == 0
|
||
|
||
|
||
def _segment_stub():
|
||
return SimpleNamespace(id=1)
|
||
|
||
|
||
def test_history_based_estimate():
|
||
events = [_event(5000, 2000), _event(7000, 3000)]
|
||
session = _PlannedSession([[_segment_stub() for _ in range(2)], events, []])
|
||
report = estimate_campaign_cost(session, _campaign())
|
||
assert report["source"] == "history"
|
||
# 均值 (6000, 2500) × 2 段 × 1.6 = 27200
|
||
assert report["estimated_tokens"] == 27200
|
||
# 费率均值 in=2.0 out=8.0：(12000/1e6*2 + 5000/1e6*8) 元 = 0.064 元 → 6 分
|
||
assert report["estimated_amount_minor"] == 6
|
||
|
||
|
||
def test_duration_from_llm_calls():
|
||
calls = [SimpleNamespace(duration_ms=10000), SimpleNamespace(duration_ms=20000)]
|
||
session = _PlannedSession([[_segment_stub()], [], calls])
|
||
report = estimate_campaign_cost(session, _campaign())
|
||
# 平均 15s × 3 + 15 = 60s 单段；1 段 + 20 = 80s
|
||
assert report["estimated_seconds"] == 80
|
||
assert report["estimated_duration_text"] == "1 分 20 秒"
|
||
|
||
|
||
def test_estimate_scales_with_segments():
|
||
one = estimate_campaign_cost(_PlannedSession([[_segment_stub()], [], []]), _campaign())
|
||
four = estimate_campaign_cost(_PlannedSession([[_segment_stub() for _ in range(4)], [], []]), _campaign())
|
||
assert four["estimated_tokens"] == one["estimated_tokens"] * 4
|
||
assert four["estimated_points"] > one["estimated_points"]
|