Aether/tests/services/billing/test_dimension_collector_service.py

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]