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,
    filter_proxy_response_headers,
)


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

    def test_new_format_only(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_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_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": 500,
        }
        assert extract_cache_creation_tokens(usage) == 500

    def test_both_formats_prefers_new(self) -> None:
        """测试同时存在时优先使用新格式"""
        usage = {
            "claude_cache_creation_5_m_tokens": 100,
            "claude_cache_creation_1_h_tokens": 200,
            "cache_creation_input_tokens": 999,  # 应该被忽略
        }
        assert extract_cache_creation_tokens(usage) == 300

    def test_empty_usage(self) -> None:
        """测试空字典"""
        usage = {}
        assert extract_cache_creation_tokens(usage) == 0

    def test_all_zeros(self) -> None:
        """测试所有字段都为 0"""
        usage = {
            "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_partial_new_format_with_old_format_fallback(self) -> None:
        """测试新格式字段不存在时回退到旧格式"""
        usage = {
            "cache_creation_input_tokens": 123,
        }
        assert extract_cache_creation_tokens(usage) == 123

    def test_new_format_zero_should_not_fallback(self) -> None:
        """测试新格式字段存在但为 0 时，不应 fallback 到旧格式"""
        usage = {
            "claude_cache_creation_5_m_tokens": 0,
            "claude_cache_creation_1_h_tokens": 0,
            "cache_creation_input_tokens": 456,
        }
        # 新格式字段存在，即使值为 0 也应该使用新格式（返回 0）
        # 而不是 fallback 到旧格式（返回 456）
        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,
            "claude_cache_creation_5_m_tokens": 50,
            "claude_cache_creation_1_h_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"


class TestFilterProxyResponseHeaders:
    def test_none_returns_empty(self) -> None:
        assert filter_proxy_response_headers(None) == {}

    def test_filters_blocklisted_headers_case_insensitive(self) -> None:
        headers = {
            "Content-Length": "123",
            "content-encoding": "gzip",
            "Transfer-Encoding": "chunked",
            "Connection": "keep-alive",
            "Keep-Alive": "timeout=5",
            "Content-Type": "application/json",
            "X-Request-Id": "abc",
            "Anthropic-RateLimit-Requests-Remaining": "10",
        }

        result = filter_proxy_response_headers(headers)

        assert "Content-Length" not in result
        assert "content-encoding" not in result
        assert "Transfer-Encoding" not in result
        assert "Connection" not in result
        assert "Keep-Alive" not in result
        assert "Content-Type" not in result

        assert result["X-Request-Id"] == "abc"
        assert result["Anthropic-RateLimit-Requests-Remaining"] == "10"