Aether/tests/api/handlers/base/test_utils.py

"""测试 handler 基础工具函数"""

import pytest

from src.api.handlers.base.utils import build_sse_headers, extract_cache_creation_tokens


class TestExtractCacheCreationTokens:
    """测试 extract_cache_creation_tokens 函数"""

    # === 嵌套格式测试（优先级最高）===

    def test_nested_cache_creation_format(self) -> None:
        """测试嵌套格式正常情况"""
        usage = {
            "cache_creation": {
                "ephemeral_5m_input_tokens": 456,
                "ephemeral_1h_input_tokens": 100,
            }
        }
        assert extract_cache_creation_tokens(usage) == 556

    def test_nested_cache_creation_with_old_format_fallback(self) -> None:
        """测试嵌套格式为 0 时回退到旧格式"""
        usage = {
            "cache_creation": {
                "ephemeral_5m_input_tokens": 0,
                "ephemeral_1h_input_tokens": 0,
            },
            "cache_creation_input_tokens": 549,
        }
        assert extract_cache_creation_tokens(usage) == 549

    def test_nested_has_priority_over_flat(self) -> None:
        """测试嵌套格式优先于扁平格式"""
        usage = {
            "cache_creation": {
                "ephemeral_5m_input_tokens": 100,
                "ephemeral_1h_input_tokens": 200,
            },
            "claude_cache_creation_5_m_tokens": 999,  # 应该被忽略
            "claude_cache_creation_1_h_tokens": 888,  # 应该被忽略
            "cache_creation_input_tokens": 777,  # 应该被忽略
        }
        assert extract_cache_creation_tokens(usage) == 300

    # === 扁平格式测试（优先级第二）===

    def test_flat_new_format_still_works(self) -> None:
        """测试扁平新格式兼容性"""
        usage = {
            "claude_cache_creation_5_m_tokens": 100,
            "claude_cache_creation_1_h_tokens": 200,
        }
        assert extract_cache_creation_tokens(usage) == 300

    def test_flat_new_format_with_old_format_fallback(self) -> None:
        """测试扁平格式为 0 时回退到旧格式"""
        usage = {
            "claude_cache_creation_5_m_tokens": 0,
            "claude_cache_creation_1_h_tokens": 0,
            "cache_creation_input_tokens": 549,
        }
        assert extract_cache_creation_tokens(usage) == 549

    def test_flat_new_format_5m_only(self) -> None:
        """测试只有 5 分钟扁平缓存"""
        usage = {
            "claude_cache_creation_5_m_tokens": 150,
            "claude_cache_creation_1_h_tokens": 0,
        }
        assert extract_cache_creation_tokens(usage) == 150

    def test_flat_new_format_1h_only(self) -> None:
        """测试只有 1 小时扁平缓存"""
        usage = {
            "claude_cache_creation_5_m_tokens": 0,
            "claude_cache_creation_1_h_tokens": 250,
        }
        assert extract_cache_creation_tokens(usage) == 250

    # === 旧格式测试（优先级第三）===

    def test_old_format_only(self) -> None:
        """测试只有旧格式"""
        usage = {
            "cache_creation_input_tokens": 549,
        }
        assert extract_cache_creation_tokens(usage) == 549

    # === 边界情况测试 ===

    def test_no_cache_creation_tokens(self) -> None:
        """测试没有任何缓存字段"""
        usage = {}
        assert extract_cache_creation_tokens(usage) == 0

    def test_all_formats_zero(self) -> None:
        """测试所有格式都为 0"""
        usage = {
            "cache_creation": {
                "ephemeral_5m_input_tokens": 0,
                "ephemeral_1h_input_tokens": 0,
            },
            "claude_cache_creation_5_m_tokens": 0,
            "claude_cache_creation_1_h_tokens": 0,
            "cache_creation_input_tokens": 0,
        }
        assert extract_cache_creation_tokens(usage) == 0

    def test_unrelated_fields_ignored(self) -> None:
        """测试忽略无关字段"""
        usage = {
            "input_tokens": 1000,
            "output_tokens": 2000,
            "cache_read_input_tokens": 300,
            "cache_creation": {
                "ephemeral_5m_input_tokens": 50,
                "ephemeral_1h_input_tokens": 75,
            },
        }
        assert extract_cache_creation_tokens(usage) == 125


class TestBuildSSEHeaders:
    def test_default_headers(self) -> None:
        headers = build_sse_headers()
        assert headers["Cache-Control"] == "no-cache, no-transform"
        assert headers["X-Accel-Buffering"] == "no"

    def test_merge_extra_headers(self) -> None:
        headers = build_sse_headers({"X-Test": "1", "Cache-Control": "custom"})
        assert headers["X-Test"] == "1"
        assert headers["Cache-Control"] == "custom"