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

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

import gzip
import json

import pytest

from src.api.handlers.base.utils import (
    build_json_response_for_client,
    build_sse_headers,
    filter_proxy_response_headers,
    resolve_client_accept_encoding,
    resolve_client_content_encoding,
)
from src.core.usage_tokens import extract_cache_creation_tokens


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: dict[str, int] = {}
        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"


class TestResolveClientEncoding:
    def test_content_encoding_prefers_hint(self) -> None:
        headers = {"content-encoding": "gzip"}
        result = resolve_client_content_encoding(headers, hinted_content_encoding="br")
        assert result == "br"

    def test_content_encoding_fallback_to_headers(self) -> None:
        headers = {"Content-Encoding": "gzip"}
        result = resolve_client_content_encoding(headers)
        assert result == "gzip"

    def test_accept_encoding_prefers_hint(self) -> None:
        headers = {"accept-encoding": "gzip"}
        result = resolve_client_accept_encoding(headers, hinted_accept_encoding="br")
        assert result == "br"

    def test_accept_encoding_fallback_to_headers(self) -> None:
        headers = {"Accept-Encoding": "gzip, deflate"}
        result = resolve_client_accept_encoding(headers)
        assert result == "gzip, deflate"


class TestBuildJsonResponseForClient:
    def test_returns_gzip_response_when_client_accepts_gzip(self) -> None:
        response = build_json_response_for_client(
            status_code=200,
            content={"ok": True},
            headers={"content-type": "application/json"},
            client_accept_encoding="gzip, deflate",
        )

        assert response.headers.get("content-encoding") == "gzip"
        assert "accept-encoding" in response.headers.get("vary", "").lower()
        decompressed = gzip.decompress(bytes(response.body))
        assert json.loads(decompressed.decode("utf-8")) == {"ok": True}

    def test_returns_plain_json_when_gzip_not_accepted(self) -> None:
        response = build_json_response_for_client(
            status_code=200,
            content={"ok": True},
            headers={"content-type": "application/json"},
            client_accept_encoding="gzip;q=0, deflate",
        )

        assert response.headers.get("content-encoding") is None
        assert json.loads(bytes(response.body).decode("utf-8")) == {"ok": True}