import httpx import astrbot.core.provider.sources.anthropic_source as anthropic_source import astrbot.core.provider.sources.kimi_code_source as kimi_code_source class _FakeAsyncAnthropic: def __init__(self, **kwargs): self.kwargs = kwargs async def close(self): return None def test_anthropic_provider_injects_custom_headers_into_http_client(monkeypatch): monkeypatch.setattr(anthropic_source, "AsyncAnthropic", _FakeAsyncAnthropic) provider = anthropic_source.ProviderAnthropic( provider_config={ "id": "anthropic-test", "type": "anthropic_chat_completion", "model": "claude-test", "key": ["test-key"], "custom_headers": { "User-Agent": "custom-agent/1.0", "X-Test-Header": 123, }, }, provider_settings={}, ) assert provider.custom_headers == { "User-Agent": "custom-agent/1.0", "X-Test-Header": "123", } assert isinstance(provider.client.kwargs["http_client"], httpx.AsyncClient) assert provider.client.kwargs["http_client"].headers["User-Agent"] == "custom-agent/1.0" assert provider.client.kwargs["http_client"].headers["X-Test-Header"] == "123" def test_kimi_code_provider_sets_defaults_and_preserves_custom_headers(monkeypatch): monkeypatch.setattr(anthropic_source, "AsyncAnthropic", _FakeAsyncAnthropic) provider = kimi_code_source.ProviderKimiCode( provider_config={ "id": "kimi-code", "type": "kimi_code_chat_completion", "key": ["test-key"], "custom_headers": {"X-Trace-Id": "trace-1"}, }, provider_settings={}, ) assert provider.base_url == kimi_code_source.KIMI_CODE_API_BASE assert provider.get_model() == kimi_code_source.KIMI_CODE_DEFAULT_MODEL assert provider.custom_headers == { "User-Agent": kimi_code_source.KIMI_CODE_USER_AGENT, "X-Trace-Id": "trace-1", } assert provider.client.kwargs["http_client"].headers["User-Agent"] == ( kimi_code_source.KIMI_CODE_USER_AGENT ) assert provider.client.kwargs["http_client"].headers["X-Trace-Id"] == "trace-1" def test_kimi_code_provider_restores_required_user_agent_when_blank(monkeypatch): monkeypatch.setattr(anthropic_source, "AsyncAnthropic", _FakeAsyncAnthropic) provider = kimi_code_source.ProviderKimiCode( provider_config={ "id": "kimi-code", "type": "kimi_code_chat_completion", "key": ["test-key"], "custom_headers": {"User-Agent": " "}, }, provider_settings={}, ) assert provider.custom_headers == { "User-Agent": kimi_code_source.KIMI_CODE_USER_AGENT, }