184 lines
6.8 KiB
Python
184 lines
6.8 KiB
Python
"""격리 PostgreSQL에서 관리자 비용 원장 SQL을 검증하는 선택형 통합 테스트."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import unittest
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from urllib.parse import urlparse
|
|
|
|
import asyncpg
|
|
|
|
from . import db
|
|
from .routes import admin as admin_routes
|
|
from .services.llm_pricing import estimate_reference_cost
|
|
|
|
|
|
TEST_DATABASE_URL = os.getenv("VIGNETTE_USAGE_TEST_DATABASE_URL", "").strip()
|
|
|
|
|
|
@unittest.skipUnless(
|
|
TEST_DATABASE_URL,
|
|
"VIGNETTE_USAGE_TEST_DATABASE_URL이 설정된 새 격리 DB에서만 실행",
|
|
)
|
|
class AdminUsagePostgresTest(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self) -> None:
|
|
parsed = urlparse(TEST_DATABASE_URL)
|
|
if parsed.path != "/vignette_usage_test":
|
|
self.fail("격리 DB 이름은 vignette_usage_test여야 한다")
|
|
if db._pool is not None:
|
|
self.fail("기존 애플리케이션 DB pool이 있는 프로세스에서는 실행할 수 없다")
|
|
|
|
bootstrap = await asyncpg.connect(TEST_DATABASE_URL)
|
|
try:
|
|
# 기존 스키마를 지우지 않는다. 이미 쓰인 DB라면 CREATE가 실패해 닫힌다.
|
|
await bootstrap.execute(
|
|
"""
|
|
CREATE SCHEMA app;
|
|
CREATE TABLE app.turns (
|
|
id BIGSERIAL PRIMARY KEY,
|
|
speaker TEXT NOT NULL,
|
|
llm_provider TEXT,
|
|
model TEXT,
|
|
tokens_in INT,
|
|
tokens_out INT,
|
|
cost_usd NUMERIC(12,6),
|
|
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
|
);
|
|
"""
|
|
)
|
|
migration = (
|
|
Path(__file__).resolve().parents[3]
|
|
/ "infra"
|
|
/ "db"
|
|
/ "init"
|
|
/ "18_admin_usage_ledger_index.sql"
|
|
).read_text(encoding="utf-8")
|
|
await bootstrap.execute(migration)
|
|
finally:
|
|
await bootstrap.close()
|
|
|
|
db._pool = await asyncpg.create_pool(
|
|
dsn=TEST_DATABASE_URL,
|
|
min_size=1,
|
|
max_size=2,
|
|
init=db._init_connection,
|
|
)
|
|
|
|
async def asyncTearDown(self) -> None:
|
|
if db._pool is not None:
|
|
await db._pool.close()
|
|
db._pool = None
|
|
|
|
async def test_usage_sql_preserves_pricing_boundaries_and_fake_exclusion(self) -> None:
|
|
observed_at = datetime.now(timezone.utc)
|
|
token_pairs = [(5_521, 210)] * 20 + [(5_530, 207)]
|
|
assert db._pool is not None
|
|
async with db._pool.acquire() as conn:
|
|
await conn.executemany(
|
|
"""
|
|
INSERT INTO app.turns (
|
|
speaker, llm_provider, model, tokens_in, tokens_out, cost_usd, created_at
|
|
) VALUES ('client', $1, $2, $3, $4, 0, $5)
|
|
""",
|
|
[
|
|
(
|
|
"agy_cli",
|
|
"gemini-3.7-flash-high",
|
|
tokens_in,
|
|
tokens_out,
|
|
observed_at,
|
|
)
|
|
for tokens_in, tokens_out in token_pairs
|
|
],
|
|
)
|
|
await conn.executemany(
|
|
"""
|
|
INSERT INTO app.turns (
|
|
speaker, llm_provider, model, tokens_in, tokens_out, cost_usd, created_at
|
|
) VALUES ('client', ' E2E ', ' FAKE-CLIENT ', 1, 1, 0, $1)
|
|
""",
|
|
[(observed_at,)] * 6,
|
|
)
|
|
|
|
expected = estimate_reference_cost(
|
|
provider="agy_cli",
|
|
model="gemini-3.7-flash-high",
|
|
tokens_in=115_950,
|
|
tokens_out=4_407,
|
|
priced_at=observed_at,
|
|
)
|
|
self.assertIsNotNone(expected)
|
|
assert expected is not None
|
|
|
|
usage = await admin_routes._usage_from_database(window_days=36_500)
|
|
|
|
self.assertEqual(usage.total_turns, 21)
|
|
self.assertEqual(usage.token_metered_turns, 21)
|
|
self.assertEqual(usage.tokens_in, 115_950)
|
|
self.assertEqual(usage.tokens_out, 4_407)
|
|
self.assertEqual(usage.cost_usd, round(expected.cost_usd, 6))
|
|
self.assertEqual(len(usage.by_provider), 1)
|
|
self.assertEqual(usage.by_provider[0].model, "gemini-3.7-flash-high")
|
|
self.assertEqual(usage.by_provider[0].cost_basis, "reference_upper_bound")
|
|
|
|
async with db._pool.acquire() as conn:
|
|
index_definition = await conn.fetchval(
|
|
"""
|
|
SELECT indexdef
|
|
FROM pg_indexes
|
|
WHERE schemaname = 'app'
|
|
AND tablename = 'turns'
|
|
AND indexname = 'idx_turns_admin_usage_created_at'
|
|
"""
|
|
)
|
|
self.assertIn("created_at DESC", index_definition or "")
|
|
self.assertIn("speaker = 'client'", index_definition or "")
|
|
|
|
async with db._pool.acquire() as conn:
|
|
await conn.executemany(
|
|
"""
|
|
INSERT INTO app.turns (
|
|
speaker, llm_provider, model, tokens_in, tokens_out, cost_usd, created_at
|
|
) VALUES ('client', 'agy_cli', 'gemini-3.1-pro-high', 150000, 10000, 0, $1)
|
|
""",
|
|
[(observed_at,), (observed_at,)],
|
|
)
|
|
|
|
tier_usage = await admin_routes._usage_from_database(window_days=36_500)
|
|
tier_row = next(
|
|
item for item in tier_usage.by_provider if item.model == "gemini-3.1-pro-high"
|
|
)
|
|
self.assertEqual(tier_row.turns, 2)
|
|
self.assertEqual(tier_row.cost_usd, 0.84)
|
|
|
|
async with db._pool.acquire() as conn:
|
|
await conn.executemany(
|
|
"""
|
|
INSERT INTO app.turns (
|
|
speaker, llm_provider, model, tokens_in, tokens_out, cost_usd, created_at
|
|
) VALUES ('client', 'agy_cli', 'gemini-3.7-flash-high', 1000000, 1000000, 0, $1)
|
|
""",
|
|
[
|
|
(datetime(2026, 8, 12, 12, tzinfo=timezone.utc),),
|
|
(datetime(2026, 8, 13, 12, tzinfo=timezone.utc),),
|
|
],
|
|
)
|
|
|
|
partial_usage = await admin_routes._usage_from_database(window_days=36_500)
|
|
partial_row = next(
|
|
item
|
|
for item in partial_usage.by_provider
|
|
if item.model == "gemini-3.7-flash-high"
|
|
)
|
|
self.assertEqual(partial_row.turns, 23)
|
|
self.assertEqual(partial_row.cost_usd, round(expected.cost_usd + 4.5, 6))
|
|
self.assertEqual(partial_usage.cost_basis, "partial_upper_bound")
|
|
self.assertEqual(partial_usage.budget.cost_basis, "partial_upper_bound")
|
|
self.assertEqual(partial_row.cost_basis, "partial_upper_bound")
|
|
self.assertIn("일부 호출 미산정", partial_row.rate_label or "")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|