mirror of
https://github.com/katanemo/plano.git
synced 2026-04-25 00:36:34 +02:00
88a02dc478
* Some fixes on model server * Remove prompt_prefilling message * Fix logging * Fix poetry issues * Improve logging and update the support for text truncation * Fix tests * Fix tests * Fix tests * Fix modelserver tests * Update modelserver tests
47 lines
1.4 KiB
Python
47 lines
1.4 KiB
Python
import pytest
|
|
import httpx
|
|
|
|
from fastapi.testclient import TestClient
|
|
from src.main import app
|
|
|
|
|
|
client = TestClient(app)
|
|
|
|
|
|
# Unit tests for the health check endpoint
|
|
@pytest.mark.asyncio
|
|
async def test_healthz():
|
|
response = client.get("/healthz")
|
|
assert response.status_code == 200
|
|
assert response.json() == {"status": "ok"}
|
|
|
|
|
|
# Unit test for the models endpoint
|
|
@pytest.mark.asyncio
|
|
async def test_models():
|
|
response = client.get("/models")
|
|
assert response.status_code == 200
|
|
assert response.json()["object"] == "list"
|
|
assert len(response.json()["data"]) > 0
|
|
|
|
|
|
# Unit test for the guardrail endpoint
|
|
@pytest.mark.asyncio
|
|
async def test_guardrail_endpoint():
|
|
request_data = {"input": "Test for jailbreak and toxicity", "task": "jailbreak"}
|
|
response = client.post("/guardrails", json=request_data)
|
|
assert response.status_code == 200
|
|
|
|
|
|
# Unit test for the function calling endpoint
|
|
@pytest.mark.asyncio
|
|
async def test_function_calling_endpoint():
|
|
async with httpx.AsyncClient(app=app, base_url="http://test") as client:
|
|
request_data = {
|
|
"messages": [{"role": "user", "content": "Hello!"}],
|
|
"model": "Arch-Function",
|
|
"tools": [],
|
|
"metadata": {"x-arch-state": "[]"},
|
|
}
|
|
response = await client.post("/function_calling", json=request_data)
|
|
assert response.status_code == 200
|