92 lines
2.9 KiB
Python
92 lines
2.9 KiB
Python
|
import pytest
|
||
|
|
import os
|
||
|
|
import tempfile
|
||
|
|
from datetime import datetime, timezone
|
||
|
|
|
||
|
|
from agent.cost_guard import CostGuard
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
async def cost_guard():
|
||
|
|
fd, db_path = tempfile.mkstemp(suffix=".db")
|
||
|
|
os.close(fd)
|
||
|
|
guard = CostGuard(db_path=db_path, daily_limit=10.0, per_task_limit=3.0)
|
||
|
|
await guard.initialize()
|
||
|
|
yield guard
|
||
|
|
await guard.close()
|
||
|
|
os.unlink(db_path)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_record_usage(cost_guard):
|
||
|
|
"""API 사용량을 기록한다."""
|
||
|
|
await cost_guard.record_usage(
|
||
|
|
task_id="task-1",
|
||
|
|
tokens_input=1000,
|
||
|
|
tokens_output=500,
|
||
|
|
)
|
||
|
|
daily = await cost_guard.get_daily_cost()
|
||
|
|
assert daily > 0
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_calculate_cost(cost_guard):
|
||
|
|
"""토큰에서 비용을 계산한다."""
|
||
|
|
cost = cost_guard.calculate_cost(tokens_input=1_000_000, tokens_output=1_000_000)
|
||
|
|
# input: $3/MTok + output: $15/MTok = $18
|
||
|
|
assert cost == pytest.approx(18.0, rel=0.01)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_check_daily_limit_ok(cost_guard):
|
||
|
|
"""일일 한도 내에서 True를 반환한다."""
|
||
|
|
result = await cost_guard.check_daily_limit()
|
||
|
|
assert result is True
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_check_daily_limit_exceeded(cost_guard):
|
||
|
|
"""일일 한도 초과 시 False를 반환한다."""
|
||
|
|
for i in range(5):
|
||
|
|
await cost_guard.record_usage(f"task-{i}", tokens_input=1_000_000, tokens_output=200_000)
|
||
|
|
result = await cost_guard.check_daily_limit()
|
||
|
|
assert result is False
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_check_task_limit_ok(cost_guard):
|
||
|
|
"""작업당 한도 내에서 True를 반환한다."""
|
||
|
|
await cost_guard.record_usage("task-1", tokens_input=100, tokens_output=50)
|
||
|
|
result = await cost_guard.check_task_limit("task-1")
|
||
|
|
assert result is True
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_check_task_limit_exceeded(cost_guard):
|
||
|
|
"""작업당 한도 초과 시 False를 반환한다."""
|
||
|
|
await cost_guard.record_usage("task-1", tokens_input=1_000_000, tokens_output=100_000)
|
||
|
|
result = await cost_guard.check_task_limit("task-1")
|
||
|
|
assert result is False
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_get_task_cost(cost_guard):
|
||
|
|
"""특정 작업의 누적 비용을 반환한다."""
|
||
|
|
await cost_guard.record_usage("task-1", tokens_input=1000, tokens_output=500)
|
||
|
|
await cost_guard.record_usage("task-1", tokens_input=2000, tokens_output=1000)
|
||
|
|
cost = await cost_guard.get_task_cost("task-1")
|
||
|
|
assert cost > 0
|
||
|
|
other_cost = await cost_guard.get_task_cost("task-2")
|
||
|
|
assert other_cost == 0.0
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_get_daily_summary(cost_guard):
|
||
|
|
"""일일 요약 정보를 반환한다."""
|
||
|
|
await cost_guard.record_usage("task-1", tokens_input=1000, tokens_output=500)
|
||
|
|
summary = await cost_guard.get_daily_summary()
|
||
|
|
assert "total_cost_usd" in summary
|
||
|
|
assert "daily_limit_usd" in summary
|
||
|
|
assert "remaining_usd" in summary
|
||
|
|
assert summary["record_count"] == 1
|