Spaces:
Sleeping
Sleeping
File size: 1,228 Bytes
bf6d237 |
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 |
from fastapi.testclient import TestClient
from private_gpt.open_ai.openai_models import OpenAICompletion, OpenAIMessage
from private_gpt.server.chat.chat_router import ChatBody
def test_chat_route_produces_a_stream(test_client: TestClient) -> None:
body = ChatBody(
messages=[OpenAIMessage(content="test", role="user")],
use_context=False,
stream=True,
)
response = test_client.post("/v1/chat/completions", json=body.model_dump())
raw_events = response.text.split("\n\n")
events = [
item.removeprefix("data: ") for item in raw_events if item.startswith("data: ")
]
assert response.status_code == 200
assert "text/event-stream" in response.headers["content-type"]
assert len(events) > 0
assert events[-1] == "[DONE]"
def test_chat_route_produces_a_single_value(test_client: TestClient) -> None:
body = ChatBody(
messages=[OpenAIMessage(content="test", role="user")],
use_context=False,
stream=False,
)
response = test_client.post("/v1/chat/completions", json=body.model_dump())
# No asserts, if it validates it's good
OpenAICompletion.model_validate(response.json())
assert response.status_code == 200
|