Spaces:
Sleeping
Sleeping
File size: 2,720 Bytes
00a2010 d86dd21 | 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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | 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_public_root_and_health_allowed_when_auth_enabled():
client = TestClient(app)
settings = Settings()
settings.anthropic_auth_token = "public-check"
app.dependency_overrides[get_settings] = lambda: settings
root_response = client.get("/")
assert root_response.status_code == 200
health_response = client.get("/health")
assert health_response.status_code == 200
assert health_response.json()["status"] == "healthy"
app.dependency_overrides.clear()
def test_api_routes_remain_protected_when_auth_enabled():
client = TestClient(app)
settings = Settings()
settings.anthropic_auth_token = "still-protected"
app.dependency_overrides[get_settings] = lambda: settings
payload = {
"model": "claude-3-sonnet",
"messages": [{"role": "user", "content": "hello"}],
}
response = client.post("/v1/messages/count_tokens", json=payload)
assert response.status_code == 401
app.dependency_overrides.clear()
|