File size: 2,193 Bytes
6172a47 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | from unittest.mock import patch
from fastapi.testclient import TestClient
from api.app import app
from api.dependencies import get_settings
from config.settings import Settings
def test_anthropic_auth_token_required_and_accepts_x_api_key():
client = TestClient(app)
settings = Settings()
settings.anthropic_auth_token = "s3cr3t"
app.dependency_overrides[get_settings] = lambda: settings
payload = {
"model": "claude-3-sonnet",
"messages": [{"role": "user", "content": "hello"}],
}
with patch("api.routes.get_token_count", return_value=1):
# No header -> 401
r = client.post("/v1/messages/count_tokens", json=payload)
assert r.status_code == 401
# X-API-Key header -> 200
r = client.post(
"/v1/messages/count_tokens", json=payload, headers={"X-API-Key": "s3cr3t"}
)
assert r.status_code == 200
assert r.json()["input_tokens"] == 1
app.dependency_overrides.clear()
def test_anthropic_auth_token_accepts_bearer_authorization():
client = TestClient(app)
settings = Settings()
settings.anthropic_auth_token = "b3artoken"
app.dependency_overrides[get_settings] = lambda: settings
payload = {
"model": "claude-3-sonnet",
"messages": [{"role": "user", "content": "hello"}],
}
with patch("api.routes.get_token_count", return_value=2):
# Authorization Bearer -> 200
r = client.post(
"/v1/messages/count_tokens",
json=payload,
headers={"Authorization": "Bearer b3artoken"},
)
assert r.status_code == 200
assert r.json()["input_tokens"] == 2
app.dependency_overrides.clear()
def test_anthropic_auth_token_applies_to_models_endpoint():
client = TestClient(app)
settings = Settings()
settings.anthropic_auth_token = "models-token"
app.dependency_overrides[get_settings] = lambda: settings
r = client.get("/v1/models")
assert r.status_code == 401
r = client.get("/v1/models", headers={"X-API-Key": "models-token"})
assert r.status_code == 200
assert "data" in r.json()
app.dependency_overrides.clear()
|