mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-15 17:30:13 +08:00
96 lines
3.2 KiB
Python
96 lines
3.2 KiB
Python
import httpx
|
|
import pytest
|
|
|
|
import astrbot.core.provider.sources.anthropic_source as anthropic_source
|
|
import astrbot.core.provider.sources.kimi_code_source as kimi_code_source
|
|
from astrbot.core.exceptions import EmptyModelOutputError
|
|
from astrbot.core.provider.entities import LLMResponse
|
|
|
|
|
|
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,
|
|
}
|
|
|
|
|
|
def test_anthropic_empty_output_raises_empty_model_output_error():
|
|
llm_response = LLMResponse(role="assistant")
|
|
|
|
with pytest.raises(EmptyModelOutputError):
|
|
anthropic_source.ProviderAnthropic._ensure_usable_response(
|
|
llm_response,
|
|
completion_id="msg_empty",
|
|
stop_reason="end_turn",
|
|
)
|