mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
123 lines
4.3 KiB
Python
123 lines
4.3 KiB
Python
from unittest.mock import MagicMock
|
|
|
|
from src.models.database import DimensionCollector
|
|
from src.services.billing.dimension_collector_service import (
|
|
DimensionCollectInput,
|
|
DimensionCollectorRuntime,
|
|
DimensionCollectorService,
|
|
)
|
|
|
|
|
|
class TestDimensionCollectorRuntime:
|
|
def test_priority_fallback(self) -> None:
|
|
runtime = DimensionCollectorRuntime()
|
|
collectors = [
|
|
DimensionCollector(
|
|
api_format="openai:chat",
|
|
task_type="chat",
|
|
dimension_name="input_tokens",
|
|
source_type="response",
|
|
source_path="usage.prompt_tokens",
|
|
value_type="int",
|
|
priority=10,
|
|
is_enabled=True,
|
|
),
|
|
DimensionCollector(
|
|
api_format="openai:chat",
|
|
task_type="chat",
|
|
dimension_name="input_tokens",
|
|
source_type="response",
|
|
source_path="usageMetadata.promptTokenCount",
|
|
value_type="int",
|
|
priority=5,
|
|
is_enabled=True,
|
|
),
|
|
]
|
|
dims = runtime.collect(
|
|
collectors=collectors, # type: ignore[arg-type]
|
|
inp=DimensionCollectInput(
|
|
response={"usageMetadata": {"promptTokenCount": 123}},
|
|
),
|
|
)
|
|
assert dims["input_tokens"] == 123
|
|
|
|
def test_transform_expression_value(self) -> None:
|
|
runtime = DimensionCollectorRuntime()
|
|
collectors = [
|
|
DimensionCollector(
|
|
api_format="gemini:video",
|
|
task_type="video",
|
|
dimension_name="file_size_mb",
|
|
source_type="metadata",
|
|
source_path="result.file_size_bytes",
|
|
transform_expression="value / 1024 / 1024",
|
|
value_type="float",
|
|
priority=0,
|
|
is_enabled=True,
|
|
)
|
|
]
|
|
dims = runtime.collect(
|
|
collectors=collectors, # type: ignore[arg-type]
|
|
inp=DimensionCollectInput(metadata={"result": {"file_size_bytes": 1048576}}),
|
|
)
|
|
assert abs(dims["file_size_mb"] - 1.0) < 1e-9
|
|
|
|
def test_computed_dimension(self) -> None:
|
|
runtime = DimensionCollectorRuntime()
|
|
collectors = [
|
|
DimensionCollector(
|
|
api_format="claude:chat",
|
|
task_type="chat",
|
|
dimension_name="input_tokens",
|
|
source_type="request",
|
|
source_path="usage.input_tokens",
|
|
value_type="int",
|
|
priority=0,
|
|
is_enabled=True,
|
|
),
|
|
DimensionCollector(
|
|
api_format="claude:chat",
|
|
task_type="chat",
|
|
dimension_name="cache_read_tokens",
|
|
source_type="request",
|
|
source_path="usage.cache_read_tokens",
|
|
value_type="int",
|
|
priority=0,
|
|
is_enabled=True,
|
|
),
|
|
DimensionCollector(
|
|
api_format="claude:chat",
|
|
task_type="chat",
|
|
dimension_name="total_input_tokens",
|
|
source_type="computed",
|
|
source_path=None,
|
|
transform_expression="input_tokens + cache_read_tokens",
|
|
value_type="int",
|
|
priority=0,
|
|
is_enabled=True,
|
|
),
|
|
]
|
|
dims = runtime.collect(
|
|
collectors=collectors, # type: ignore[arg-type]
|
|
inp=DimensionCollectInput(
|
|
request={"usage": {"input_tokens": 100, "cache_read_tokens": 20}}
|
|
),
|
|
)
|
|
assert dims["input_tokens"] == 100
|
|
assert dims["cache_read_tokens"] == 20
|
|
assert dims["total_input_tokens"] == 120
|
|
|
|
|
|
class TestDimensionCollectorService:
|
|
def test_video_fallback_merges_base_collectors(self) -> None:
|
|
from src.services.billing.cache import BillingCache
|
|
|
|
BillingCache.invalidate_all()
|
|
|
|
# code-only: openai:video should fall back to openai:chat video collectors shipped in code
|
|
db = MagicMock()
|
|
svc = DimensionCollectorService(db)
|
|
result = svc.list_enabled_collectors(api_format="openai:video", task_type="video")
|
|
|
|
assert "video_size_bytes" in [c.dimension_name for c in result]
|