ai_pipeline/tests/pipeline/api/test_rag.py

45 lines
1.3 KiB
Python

from asgi_lifespan import LifespanManager
from httpx import AsyncClient, ASGITransport
from src.pipeline.main import app
from src.pipeline.utils import logger
import pytest
import pytest_asyncio
import json
# ----------------------
# async
# ----------------------
@pytest_asyncio.fixture
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_async.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())