Isolate cloud model access before enabling product RAG workflows
Some checks failed
verify / verify (push) Has been cancelled
Some checks failed
verify / verify (push) Has been cancelled
The API and ingestion tools now use a fixed internal model gateway while governed profiles, embedding cache assignments, traceable citations, and stable API errors establish the boundaries required by later workflows. Constraint: The current Alibaba Cloud workspace rejects all three live model calls with authentication failures Rejected: Give the API or seed tools the Bailian key and direct egress | combines database access, cloud credentials, and public network access Rejected: Mix offline and Bailian vectors in one demo namespace | makes profile activation and retrieval ambiguous Confidence: high Scope-risk: moderate Reversibility: clean Directive: Keep Bailian credentials and egress exclusive to model-gateway and create a new immutable profile hash for any embedding identity change Tested: make verify; 121 backend tests; 14 frontend tests; fresh and populated Alembic upgrade-downgrade-upgrade; two idempotent offline seeds; Docker health and HTTP retrieval; isolated provider smoke Not-tested: Successful live Bailian responses because the supplied workspace credential currently fails authentication
This commit is contained in:
748
backend/tests/unit/test_model_gateway.py
Normal file
748
backend/tests/unit/test_model_gateway.py
Normal file
@@ -0,0 +1,748 @@
|
||||
import asyncio
|
||||
import json
|
||||
from collections.abc import AsyncIterator, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from typing import cast
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from starlette.requests import ClientDisconnect
|
||||
from starlette.types import Message, Scope
|
||||
|
||||
from app.core.config import Settings
|
||||
from app.model_gateway import create_model_gateway_app
|
||||
from app.ports.model_providers import (
|
||||
ChatCompletionResult,
|
||||
ChatMessage,
|
||||
ChatStreamEvent,
|
||||
EmbeddingResult,
|
||||
ModelProviderError,
|
||||
ProviderErrorKind,
|
||||
ProviderUsage,
|
||||
RankedItem,
|
||||
RerankResult,
|
||||
)
|
||||
|
||||
API_TOKEN = "test-only-api-token" # noqa: S105
|
||||
WORKER_TOKEN = "test-only-worker-token" # noqa: S105
|
||||
ALLOWED_TOKENS = {"api": API_TOKEN, "worker": WORKER_TOKEN}
|
||||
|
||||
|
||||
class _FakeEmbedding:
|
||||
def __init__(self) -> None:
|
||||
self.document_calls = 0
|
||||
self.query_calls = 0
|
||||
self.active = 0
|
||||
self.max_active = 0
|
||||
self.delay = 0.0
|
||||
self.error: Exception | None = None
|
||||
|
||||
async def _result(self, texts: Sequence[str]) -> EmbeddingResult:
|
||||
if self.error is not None:
|
||||
raise self.error
|
||||
self.active += 1
|
||||
self.max_active = max(self.max_active, self.active)
|
||||
try:
|
||||
if self.delay:
|
||||
await asyncio.sleep(self.delay)
|
||||
return EmbeddingResult(
|
||||
vectors=tuple((float(index + 1), 0.5) for index, _ in enumerate(texts)),
|
||||
model="fake-embedding",
|
||||
request_id="req-embedding",
|
||||
usage=ProviderUsage(input_tokens=len(texts), total_tokens=len(texts)),
|
||||
elapsed_ms=1.25,
|
||||
)
|
||||
finally:
|
||||
self.active -= 1
|
||||
|
||||
async def embed_documents(self, texts: Sequence[str]) -> EmbeddingResult:
|
||||
self.document_calls += 1
|
||||
return await self._result(texts)
|
||||
|
||||
async def embed_query(self, text: str) -> EmbeddingResult:
|
||||
self.query_calls += 1
|
||||
return await self._result([text])
|
||||
|
||||
|
||||
class _FakeReranker:
|
||||
def __init__(self) -> None:
|
||||
self.calls = 0
|
||||
|
||||
async def rerank(
|
||||
self,
|
||||
query: str,
|
||||
documents: Sequence[str],
|
||||
*,
|
||||
top_n: int,
|
||||
instruct: str | None = None,
|
||||
) -> RerankResult:
|
||||
del query, instruct
|
||||
self.calls += 1
|
||||
return RerankResult(
|
||||
items=tuple(
|
||||
RankedItem(index=index, relevance_score=1.0 - index / 10, document=document)
|
||||
for index, document in enumerate(documents[:top_n])
|
||||
),
|
||||
model="fake-reranker",
|
||||
request_id="req-rerank",
|
||||
usage=ProviderUsage(input_tokens=len(documents), total_tokens=len(documents)),
|
||||
elapsed_ms=2.5,
|
||||
)
|
||||
|
||||
|
||||
class _FakeChat:
|
||||
def __init__(self) -> None:
|
||||
self.complete_calls = 0
|
||||
self.stream_error: ModelProviderError | None = None
|
||||
self.stream_events: tuple[ChatStreamEvent, ...] | None = None
|
||||
|
||||
async def complete(
|
||||
self,
|
||||
messages: Sequence[ChatMessage],
|
||||
*,
|
||||
max_tokens: int,
|
||||
) -> ChatCompletionResult:
|
||||
del max_tokens
|
||||
self.complete_calls += 1
|
||||
return ChatCompletionResult(
|
||||
content=f"answer:{messages[-1].content}",
|
||||
finish_reason="stop",
|
||||
model="fake-chat",
|
||||
request_id="req-chat",
|
||||
usage=ProviderUsage(input_tokens=3, output_tokens=2, total_tokens=5),
|
||||
elapsed_ms=3.75,
|
||||
)
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
messages: Sequence[ChatMessage],
|
||||
*,
|
||||
max_tokens: int,
|
||||
) -> AsyncIterator[ChatStreamEvent]:
|
||||
del messages, max_tokens
|
||||
if self.stream_error is not None:
|
||||
raise self.stream_error
|
||||
if self.stream_events is not None:
|
||||
for event in self.stream_events:
|
||||
yield event
|
||||
return
|
||||
yield ChatStreamEvent(
|
||||
delta="第一段",
|
||||
finish_reason=None,
|
||||
model="fake-chat",
|
||||
request_id="req-stream",
|
||||
usage=ProviderUsage(),
|
||||
elapsed_ms=1.0,
|
||||
)
|
||||
yield ChatStreamEvent(
|
||||
delta="第二段",
|
||||
finish_reason="stop",
|
||||
model="fake-chat",
|
||||
request_id="req-stream",
|
||||
usage=ProviderUsage(input_tokens=3, output_tokens=2, total_tokens=5),
|
||||
elapsed_ms=2.0,
|
||||
)
|
||||
|
||||
|
||||
class _TrackedChatStream:
|
||||
def __init__(self) -> None:
|
||||
self.closed = False
|
||||
self.emitted = False
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[ChatStreamEvent]:
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> ChatStreamEvent:
|
||||
if self.emitted:
|
||||
await asyncio.Event().wait()
|
||||
raise StopAsyncIteration
|
||||
self.emitted = True
|
||||
return ChatStreamEvent(
|
||||
delta="first",
|
||||
finish_reason=None,
|
||||
model="fake-chat",
|
||||
request_id="req-stream",
|
||||
usage=ProviderUsage(),
|
||||
elapsed_ms=1.0,
|
||||
)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
class _TrackedChat(_FakeChat):
|
||||
def __init__(self, stream: _TrackedChatStream) -> None:
|
||||
super().__init__()
|
||||
self.tracked_stream = stream
|
||||
|
||||
def stream(
|
||||
self,
|
||||
messages: Sequence[ChatMessage],
|
||||
*,
|
||||
max_tokens: int,
|
||||
) -> AsyncIterator[ChatStreamEvent]:
|
||||
del messages, max_tokens
|
||||
return self.tracked_stream
|
||||
|
||||
|
||||
class _ExplodingEmbeddingResult:
|
||||
@property
|
||||
def vectors(self) -> tuple[tuple[float, ...], ...]:
|
||||
raise RuntimeError("private-dto-exception-text")
|
||||
|
||||
|
||||
class _MalformedEmbedding(_FakeEmbedding):
|
||||
async def embed_query(self, text: str) -> EmbeddingResult:
|
||||
del text
|
||||
self.query_calls += 1
|
||||
return cast(EmbeddingResult, _ExplodingEmbeddingResult())
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _client(
|
||||
*,
|
||||
embedding: _FakeEmbedding | None = None,
|
||||
reranker: _FakeReranker | None = None,
|
||||
chat: _FakeChat | None = None,
|
||||
max_concurrency: int = 4,
|
||||
) -> AsyncIterator[tuple[httpx.AsyncClient, FastAPI, _FakeEmbedding, _FakeReranker, _FakeChat]]:
|
||||
embedding = embedding or _FakeEmbedding()
|
||||
reranker = reranker or _FakeReranker()
|
||||
chat = chat or _FakeChat()
|
||||
app = create_model_gateway_app(
|
||||
embedding_provider=embedding,
|
||||
reranker=reranker,
|
||||
chat_provider=chat,
|
||||
allowed_tokens=ALLOWED_TOKENS,
|
||||
max_concurrency=max_concurrency,
|
||||
)
|
||||
async with app.router.lifespan_context(app):
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport, base_url="http://model-gateway"
|
||||
) as client:
|
||||
yield client, app, embedding, reranker, chat
|
||||
|
||||
|
||||
def _headers(caller: str = "api", token: str = API_TOKEN) -> dict[str, str]:
|
||||
return {"Authorization": f"Bearer {token}", "X-RAG-Caller": caller}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_is_unauthenticated_and_schema_endpoints_are_disabled() -> None:
|
||||
async with _client() as (client, app, _, __, ___):
|
||||
live = await client.get("/health/live")
|
||||
ready = await client.get("/health/ready")
|
||||
docs = await client.get("/docs")
|
||||
openapi = await client.get("/openapi.json")
|
||||
|
||||
assert app.docs_url is None
|
||||
assert app.redoc_url is None
|
||||
assert app.openapi_url is None
|
||||
assert live.status_code == 200
|
||||
assert ready.json() == {"status": "ready", "checks": {"configuration": "ok"}}
|
||||
assert docs.status_code == 404
|
||||
assert openapi.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_internal_authentication_binds_token_to_declared_caller() -> None:
|
||||
async with _client() as (client, _, embedding, __, ___):
|
||||
payload = {"texts": ["斑岩铜矿"], "input_type": "query"}
|
||||
missing = await client.post("/internal/v1/embeddings", json=payload)
|
||||
wrong = await client.post(
|
||||
"/internal/v1/embeddings",
|
||||
json=payload,
|
||||
headers=_headers(token="wrong-and-sensitive-token"),
|
||||
)
|
||||
mismatched = await client.post(
|
||||
"/internal/v1/embeddings",
|
||||
json=payload,
|
||||
headers=_headers(caller="worker", token=API_TOKEN),
|
||||
)
|
||||
|
||||
assert [missing.status_code, wrong.status_code, mismatched.status_code] == [401, 401, 401]
|
||||
assert missing.json()["error"]["kind"] == "unauthorized"
|
||||
assert "wrong-and-sensitive-token" not in wrong.text
|
||||
assert embedding.query_calls == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_embedding_scope_and_response_contract() -> None:
|
||||
async with _client() as (client, _, embedding, __, ___):
|
||||
query = await client.post(
|
||||
"/internal/v1/embeddings",
|
||||
json={"texts": ["斑岩铜矿"], "input_type": "query"},
|
||||
headers=_headers(),
|
||||
)
|
||||
forbidden = await client.post(
|
||||
"/internal/v1/embeddings",
|
||||
json={"texts": ["内部地质报告"], "input_type": "document"},
|
||||
headers=_headers(),
|
||||
)
|
||||
document = await client.post(
|
||||
"/internal/v1/embeddings",
|
||||
json={"texts": ["文档一", "文档二"], "input_type": "document"},
|
||||
headers=_headers(caller="worker", token=WORKER_TOKEN),
|
||||
)
|
||||
|
||||
assert query.status_code == 200
|
||||
assert query.json() == {
|
||||
"vectors": [[1.0, 0.5]],
|
||||
"model": "fake-embedding",
|
||||
"request_id": "req-embedding",
|
||||
"usage": {"input_tokens": 1, "output_tokens": None, "total_tokens": 1},
|
||||
"elapsed_ms": 1.25,
|
||||
}
|
||||
assert forbidden.status_code == 403
|
||||
assert document.status_code == 200
|
||||
assert len(document.json()["vectors"]) == 2
|
||||
assert embedding.query_calls == 1
|
||||
assert embedding.document_calls == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validation_error_does_not_echo_rejected_content() -> None:
|
||||
sensitive = "private-geological-report-fragment"
|
||||
async with _client() as (client, _, embedding, __, ___):
|
||||
response = await client.post(
|
||||
"/internal/v1/embeddings",
|
||||
json={"texts": [sensitive, "second query"], "input_type": "query"},
|
||||
headers=_headers(),
|
||||
)
|
||||
|
||||
assert response.status_code == 422
|
||||
assert response.json() == {
|
||||
"error": {"kind": "invalid_request", "retryable": False, "request_id": None}
|
||||
}
|
||||
assert sensitive not in response.text
|
||||
assert embedding.query_calls == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rerank_and_chat_completion_contracts() -> None:
|
||||
async with _client() as (client, _, __, reranker, chat):
|
||||
rerank_response = await client.post(
|
||||
"/internal/v1/rerank",
|
||||
json={"query": "铜矿", "documents": ["铜矿蚀变", "煤层"], "top_n": 1},
|
||||
headers=_headers(),
|
||||
)
|
||||
chat_response = await client.post(
|
||||
"/internal/v1/chat/completions",
|
||||
json={"messages": [{"role": "user", "content": "结论是什么?"}], "max_tokens": 32},
|
||||
headers=_headers(),
|
||||
)
|
||||
|
||||
assert rerank_response.status_code == 200
|
||||
assert rerank_response.json()["items"] == [
|
||||
{"index": 0, "relevance_score": 1.0, "document": "铜矿蚀变"}
|
||||
]
|
||||
assert chat_response.status_code == 200
|
||||
assert chat_response.json()["content"] == "answer:结论是什么?"
|
||||
assert chat_response.json()["usage"]["total_tokens"] == 5
|
||||
assert reranker.calls == 1
|
||||
assert chat.complete_calls == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_error_mapping_is_stable_and_redacted() -> None:
|
||||
embedding = _FakeEmbedding()
|
||||
embedding.error = ModelProviderError(
|
||||
operation="embedding.create.private-query",
|
||||
kind=ProviderErrorKind.RATE_LIMITED,
|
||||
status_code=429,
|
||||
provider_code="private-upstream-body",
|
||||
request_id="req-safe",
|
||||
retryable=True,
|
||||
)
|
||||
async with _client(embedding=embedding) as (client, _, __, ___, ____):
|
||||
response = await client.post(
|
||||
"/internal/v1/embeddings",
|
||||
json={"texts": ["private-input-text"], "input_type": "query"},
|
||||
headers=_headers(),
|
||||
)
|
||||
|
||||
assert response.status_code == 429
|
||||
assert response.json() == {
|
||||
"error": {"kind": "rate_limited", "retryable": True, "request_id": "req-safe"}
|
||||
}
|
||||
assert "private" not in response.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_provider_and_dto_failures_are_locally_sanitized(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
provider_failure = _FakeEmbedding()
|
||||
provider_failure.error = RuntimeError("private-provider-exception-text")
|
||||
async with _client(embedding=provider_failure) as (client, _, __, ___, ____):
|
||||
provider_response = await client.post(
|
||||
"/internal/v1/embeddings",
|
||||
json={"texts": ["private-provider-input"], "input_type": "query"},
|
||||
headers=_headers(),
|
||||
)
|
||||
|
||||
malformed = _MalformedEmbedding()
|
||||
async with _client(embedding=malformed) as (client, _, __, ___, ____):
|
||||
dto_response = await client.post(
|
||||
"/internal/v1/embeddings",
|
||||
json={"texts": ["private-dto-input"], "input_type": "query"},
|
||||
headers=_headers(),
|
||||
)
|
||||
|
||||
safe_error = {"error": {"kind": "invalid_response", "retryable": False, "request_id": None}}
|
||||
assert provider_response.status_code == 502
|
||||
assert provider_response.json() == safe_error
|
||||
assert dto_response.status_code == 502
|
||||
assert dto_response.json() == safe_error
|
||||
assert "private-provider" not in provider_response.text
|
||||
assert "private-dto" not in dto_response.text
|
||||
assert "private-provider-exception-text" not in caplog.text
|
||||
assert "private-dto-exception-text" not in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sse_emits_safe_delta_complete_and_error_events() -> None:
|
||||
chat = _FakeChat()
|
||||
async with _client(chat=chat) as (client, _, __, ___, ____):
|
||||
success = await client.post(
|
||||
"/internal/v1/chat/stream",
|
||||
json={"messages": [{"role": "user", "content": "回答"}], "max_tokens": 32},
|
||||
headers=_headers(),
|
||||
)
|
||||
|
||||
assert success.status_code == 200
|
||||
assert success.headers["content-type"].startswith("text/event-stream")
|
||||
assert success.text.count("event: delta") == 2
|
||||
assert "event: complete" in success.text
|
||||
assert '"total_tokens":5' in success.text
|
||||
|
||||
chat.stream_error = ModelProviderError(
|
||||
operation="chat.stream.private",
|
||||
kind=ProviderErrorKind.AUTHENTICATION,
|
||||
provider_code="private-secret-body",
|
||||
request_id="req-stream-safe",
|
||||
)
|
||||
async with _client(chat=chat) as (client, _, __, ___, ____):
|
||||
failure = await client.post(
|
||||
"/internal/v1/chat/stream",
|
||||
json={"messages": [{"role": "user", "content": "private-question"}]},
|
||||
headers=_headers(),
|
||||
)
|
||||
|
||||
assert failure.status_code == 200
|
||||
assert "event: error" in failure.text
|
||||
assert '"kind":"authentication"' in failure.text
|
||||
assert "private" not in failure.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sse_aggregates_terminal_and_later_usage_only_event() -> None:
|
||||
chat = _FakeChat()
|
||||
chat.stream_events = (
|
||||
ChatStreamEvent(
|
||||
delta="答案",
|
||||
finish_reason=None,
|
||||
model="fake-chat-first",
|
||||
request_id="req-first",
|
||||
usage=ProviderUsage(input_tokens=3),
|
||||
elapsed_ms=4.0,
|
||||
),
|
||||
ChatStreamEvent(
|
||||
delta="",
|
||||
finish_reason="stop",
|
||||
model="",
|
||||
request_id=None,
|
||||
usage=ProviderUsage(),
|
||||
elapsed_ms=2.0,
|
||||
),
|
||||
ChatStreamEvent(
|
||||
delta="",
|
||||
finish_reason=None,
|
||||
model="fake-chat-final",
|
||||
request_id="req-final",
|
||||
usage=ProviderUsage(output_tokens=2, total_tokens=5),
|
||||
elapsed_ms=7.0,
|
||||
),
|
||||
)
|
||||
|
||||
async with _client(chat=chat) as (client, _, __, ___, ____):
|
||||
response = await client.post(
|
||||
"/internal/v1/chat/stream",
|
||||
json={"messages": [{"role": "user", "content": "回答"}]},
|
||||
headers=_headers(),
|
||||
)
|
||||
|
||||
complete_block = next(
|
||||
block for block in response.text.split("\n\n") if block.startswith("event: complete")
|
||||
)
|
||||
complete_payload = json.loads(complete_block.split("data: ", maxsplit=1)[1])
|
||||
assert complete_payload == {
|
||||
"elapsed_ms": 7.0,
|
||||
"finish_reason": "stop",
|
||||
"model": "fake-chat-final",
|
||||
"request_id": "req-final",
|
||||
"usage": {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sse_without_legal_terminal_finish_emits_error_not_complete() -> None:
|
||||
chat = _FakeChat()
|
||||
chat.stream_events = (
|
||||
ChatStreamEvent(
|
||||
delta="未完成答案",
|
||||
finish_reason=None,
|
||||
model="fake-chat",
|
||||
request_id="req-incomplete",
|
||||
usage=ProviderUsage(),
|
||||
elapsed_ms=1.0,
|
||||
),
|
||||
ChatStreamEvent(
|
||||
delta="",
|
||||
finish_reason=None,
|
||||
model="fake-chat",
|
||||
request_id="req-incomplete",
|
||||
usage=ProviderUsage(total_tokens=5),
|
||||
elapsed_ms=2.0,
|
||||
),
|
||||
)
|
||||
|
||||
async with _client(chat=chat) as (client, _, __, ___, ____):
|
||||
response = await client.post(
|
||||
"/internal/v1/chat/stream",
|
||||
json={"messages": [{"role": "user", "content": "回答"}]},
|
||||
headers=_headers(),
|
||||
)
|
||||
|
||||
assert "event: complete" not in response.text
|
||||
assert "event: error" in response.text
|
||||
assert '"kind":"invalid_response"' in response.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrency_is_bounded_across_requests() -> None:
|
||||
embedding = _FakeEmbedding()
|
||||
embedding.delay = 0.02
|
||||
async with _client(embedding=embedding, max_concurrency=1) as (client, _, __, ___, ____):
|
||||
responses = await asyncio.gather(
|
||||
client.post(
|
||||
"/internal/v1/embeddings",
|
||||
json={"texts": ["query one"], "input_type": "query"},
|
||||
headers=_headers(),
|
||||
),
|
||||
client.post(
|
||||
"/internal/v1/embeddings",
|
||||
json={"texts": ["query two"], "input_type": "query"},
|
||||
headers=_headers(),
|
||||
),
|
||||
)
|
||||
|
||||
assert [response.status_code for response in responses] == [200, 200]
|
||||
assert embedding.max_active == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_is_closed_when_downstream_disconnects() -> None:
|
||||
tracked_stream = _TrackedChatStream()
|
||||
chat = _TrackedChat(tracked_stream)
|
||||
embedding = _FakeEmbedding()
|
||||
reranker = _FakeReranker()
|
||||
app = create_model_gateway_app(
|
||||
embedding_provider=embedding,
|
||||
reranker=reranker,
|
||||
chat_provider=chat,
|
||||
allowed_tokens=ALLOWED_TOKENS,
|
||||
)
|
||||
body = json.dumps(
|
||||
{"messages": [{"role": "user", "content": "stream"}], "max_tokens": 8}
|
||||
).encode()
|
||||
request_delivered = False
|
||||
|
||||
async def receive() -> Message:
|
||||
nonlocal request_delivered
|
||||
if not request_delivered:
|
||||
request_delivered = True
|
||||
return {"type": "http.request", "body": body, "more_body": False}
|
||||
return {"type": "http.disconnect"}
|
||||
|
||||
async def disconnect_on_first_body(message: Message) -> None:
|
||||
if message["type"] == "http.response.body" and message.get("body"):
|
||||
raise OSError("downstream disconnected")
|
||||
|
||||
scope = cast(
|
||||
Scope,
|
||||
{
|
||||
"type": "http",
|
||||
"asgi": {"version": "3.0", "spec_version": "2.4"},
|
||||
"http_version": "1.1",
|
||||
"method": "POST",
|
||||
"scheme": "http",
|
||||
"path": "/internal/v1/chat/stream",
|
||||
"raw_path": b"/internal/v1/chat/stream",
|
||||
"query_string": b"",
|
||||
"headers": [
|
||||
(b"content-type", b"application/json"),
|
||||
(b"content-length", str(len(body)).encode()),
|
||||
(b"authorization", f"Bearer {API_TOKEN}".encode()),
|
||||
(b"x-rag-caller", b"api"),
|
||||
],
|
||||
"client": ("127.0.0.1", 12345),
|
||||
"server": ("model-gateway", 8000),
|
||||
},
|
||||
)
|
||||
|
||||
async with app.router.lifespan_context(app):
|
||||
with pytest.raises(ClientDisconnect):
|
||||
await app(scope, receive, disconnect_on_first_body)
|
||||
|
||||
assert tracked_stream.emitted is True
|
||||
assert tracked_stream.closed is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_production_readiness_reads_local_secrets_without_cloud_call(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
api_key = tmp_path / "bailian"
|
||||
api_token = tmp_path / "api-token"
|
||||
worker_token = tmp_path / "worker-token"
|
||||
api_key.write_text("test-only-bailian-key", encoding="utf-8")
|
||||
api_token.write_text(API_TOKEN, encoding="utf-8")
|
||||
worker_token.write_text(WORKER_TOKEN, encoding="utf-8")
|
||||
monkeypatch.setenv(
|
||||
"MODEL_GATEWAY_ALLOWED_TOKEN_FILES",
|
||||
f"api={api_token},worker={worker_token}",
|
||||
)
|
||||
settings = Settings(
|
||||
dashscope_api_key_file=api_key,
|
||||
bailian_openai_base_url=(
|
||||
"https://workspace-test.cn-beijing.maas.aliyuncs.com/compatible-mode/v1"
|
||||
),
|
||||
bailian_rerank_base_url=(
|
||||
"https://workspace-test.cn-beijing.maas.aliyuncs.com/compatible-api/v1"
|
||||
),
|
||||
)
|
||||
app = create_model_gateway_app(settings_factory=lambda: settings)
|
||||
|
||||
async with app.router.lifespan_context(app):
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport, base_url="http://model-gateway"
|
||||
) as client:
|
||||
response = await client.get("/health/ready")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"status": "ready", "checks": {"configuration": "ok"}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rotated_internal_token_fuses_old_process_until_coordinated_restart(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
api_key = tmp_path / "bailian"
|
||||
api_token = tmp_path / "api-token"
|
||||
worker_token = tmp_path / "worker-token"
|
||||
api_key.write_text("test-only-bailian-key", encoding="utf-8")
|
||||
api_token.write_text(API_TOKEN, encoding="utf-8")
|
||||
worker_token.write_text(WORKER_TOKEN, encoding="utf-8")
|
||||
monkeypatch.setenv(
|
||||
"MODEL_GATEWAY_ALLOWED_TOKEN_FILES",
|
||||
f"api={api_token},worker={worker_token}",
|
||||
)
|
||||
settings = Settings(
|
||||
dashscope_api_key_file=api_key,
|
||||
bailian_openai_base_url=(
|
||||
"https://workspace-test.cn-beijing.maas.aliyuncs.com/compatible-mode/v1"
|
||||
),
|
||||
bailian_rerank_base_url=(
|
||||
"https://workspace-test.cn-beijing.maas.aliyuncs.com/compatible-api/v1"
|
||||
),
|
||||
)
|
||||
app = create_model_gateway_app(settings_factory=lambda: settings)
|
||||
|
||||
async with app.router.lifespan_context(app):
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport, base_url="http://model-gateway"
|
||||
) as client:
|
||||
assert (await client.get("/health/ready")).status_code == 200
|
||||
rotated_token = "test-only-rotated-api-token"
|
||||
api_token.write_text(rotated_token, encoding="utf-8")
|
||||
old_credential = await client.post(
|
||||
"/internal/v1/embeddings",
|
||||
json={"texts": ["query"], "input_type": "query"},
|
||||
headers=_headers(),
|
||||
)
|
||||
new_credential = await client.post(
|
||||
"/internal/v1/embeddings",
|
||||
json={"texts": ["query"], "input_type": "query"},
|
||||
headers=_headers(token=rotated_token),
|
||||
)
|
||||
unhealthy = await client.get("/health/ready")
|
||||
|
||||
expected_restart = {
|
||||
"error": {"kind": "restart_required", "retryable": False, "request_id": None}
|
||||
}
|
||||
assert old_credential.status_code == 503
|
||||
assert old_credential.json() == expected_restart
|
||||
assert new_credential.status_code == 503
|
||||
assert new_credential.json() == expected_restart
|
||||
assert unhealthy.status_code == 503
|
||||
assert unhealthy.json()["checks"] == {"configuration": "restart_required"}
|
||||
|
||||
replacement = create_model_gateway_app(settings_factory=lambda: settings)
|
||||
async with replacement.router.lifespan_context(replacement):
|
||||
transport = httpx.ASGITransport(app=replacement)
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport, base_url="http://replacement-model-gateway"
|
||||
) as client:
|
||||
recovered = await client.get("/health/ready")
|
||||
assert recovered.status_code == 200
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rotated_bailian_key_is_detected_before_any_business_call(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
api_key = tmp_path / "bailian"
|
||||
api_token = tmp_path / "api-token"
|
||||
worker_token = tmp_path / "worker-token"
|
||||
api_key.write_text("test-only-bailian-key", encoding="utf-8")
|
||||
api_token.write_text(API_TOKEN, encoding="utf-8")
|
||||
worker_token.write_text(WORKER_TOKEN, encoding="utf-8")
|
||||
monkeypatch.setenv(
|
||||
"MODEL_GATEWAY_ALLOWED_TOKEN_FILES",
|
||||
f"api={api_token},worker={worker_token}",
|
||||
)
|
||||
settings = Settings(
|
||||
dashscope_api_key_file=api_key,
|
||||
bailian_openai_base_url=(
|
||||
"https://workspace-test.cn-beijing.maas.aliyuncs.com/compatible-mode/v1"
|
||||
),
|
||||
bailian_rerank_base_url=(
|
||||
"https://workspace-test.cn-beijing.maas.aliyuncs.com/compatible-api/v1"
|
||||
),
|
||||
)
|
||||
app = create_model_gateway_app(settings_factory=lambda: settings)
|
||||
|
||||
async with app.router.lifespan_context(app):
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport, base_url="http://model-gateway"
|
||||
) as client:
|
||||
assert (await client.get("/health/ready")).status_code == 200
|
||||
api_key.write_text("test-only-rotated-bailian-key", encoding="utf-8")
|
||||
response = await client.post(
|
||||
"/internal/v1/embeddings",
|
||||
json={"texts": ["must-not-reach-cloud"], "input_type": "query"},
|
||||
headers=_headers(),
|
||||
)
|
||||
|
||||
assert response.status_code == 503
|
||||
assert response.json()["error"]["kind"] == "restart_required"
|
||||
Reference in New Issue
Block a user