63 lines
1.7 KiB
Python
63 lines
1.7 KiB
Python
from asgi_lifespan import LifespanManager
|
|
from httpx import AsyncClient, ASGITransport
|
|
from src.pipeline.main import app
|
|
from src.pipeline.utils import logger
|
|
from fastapi.testclient import TestClient
|
|
import pytest
|
|
import pytest_asyncio
|
|
import json
|
|
|
|
# ----------------------
|
|
# async
|
|
# ----------------------
|
|
|
|
|
|
@pytest_asyncio.fixture(scope="session")
|
|
async def app_client_async():
|
|
async with LifespanManager(app):
|
|
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
|
yield client
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_chat(app_client_async):
|
|
res = await app_client.post("/api/v1/rag/stream-chat", json={"query": "hello"})
|
|
assert res.status_code == 200
|
|
sse_messages = []
|
|
async for line in res.aiter_lines():
|
|
if not line :
|
|
continue
|
|
if line.startswith("data:"):
|
|
payload = line[len("data:") :].strip()
|
|
try:
|
|
obj = json.loads(payload)
|
|
except Exception:
|
|
obj = {"raw": payload}
|
|
logger.debug(obj)
|
|
sse_messages.append(obj)
|
|
assert sse_messages
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_message_list_async(app_client_async):
|
|
res = await app_client_async.get("/api/v1/rag/message-list/a1c9108c-9201-4e60-a436-505edec3f47e")
|
|
assert res.status_code == 200
|
|
logger.debug(res.json())
|
|
|
|
|
|
# ----------------------
|
|
# normal
|
|
# ----------------------
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def app_client():
|
|
client = TestClient(app=app)
|
|
return client
|
|
|
|
|
|
def test_chat_message_list(app_client):
|
|
res = app_client.get("/api/v1/rag/message-list/a1c9108c-9201-4e60-a436-505edec3f47e")
|
|
assert res.status_code == 200
|
|
logger.debug(res.json())
|