apunkt/plano
1
0
Fork 0
mirror of https://github.com/katanemo/plano.git synced 2026-05-03 04:42:49 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix precommit

Browse source
This commit is contained in:
cotran 2024-10-16 16:49:44 -07:00
parent 953b1b056b
commit b5d2f19d9b
4 changed files with 119 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

14
model_server/app/tests/test_app.py
View file

@ -2,12 +2,20 @@ import pytest
import httpx
from fastapi.testclient import TestClient
from app.main import app # Assuming your FastAPI app is in main.py
from unittest.mock import patch
import app.commons.globals as glb
import logging
logger = logging.getLogger(__name__)
client = TestClient(app)
logger.info(f"Model will be loaded on device: {glb.DEVICE}")
# Unit tests for the health check endpoint
@pytest.mark.asyncio
@patch("app.loader.glb.DEVICE", glb.DEVICE) # Mock the device to 'cpu'
async def test_healthz():
response = client.get("/healthz")
assert response.status_code == 200
@ -16,6 +24,7 @@ async def test_healthz():
# Unit test for the models endpoint
@pytest.mark.asyncio
@patch("app.loader.glb.DEVICE", glb.DEVICE) # Mock the device to 'cpu'
async def test_models():
response = client.get("/models")
assert response.status_code == 200
@ -25,6 +34,7 @@ async def test_models():
# Unit test for embeddings endpoint
@pytest.mark.asyncio
@patch("app.loader.glb.DEVICE", glb.DEVICE) # Mock the device to 'cpu'
async def test_embedding():
request_data = {"input": "Test embedding", "model": "katanemo/bge-large-en-v1.5"}
response = client.post("/embeddings", json=request_data)
@ -38,6 +48,7 @@ async def test_embedding():
# Unit test for the guard endpoint
@pytest.mark.asyncio
@patch("app.loader.glb.DEVICE", glb.DEVICE) # Mock the device to 'cpu'
async def test_guard():
request_data = {"input": "Test for jailbreak and toxicity", "task": "jailbreak"}
response = client.post("/guard", json=request_data)
@ -47,6 +58,7 @@ async def test_guard():
# Unit test for the zero-shot endpoint
@pytest.mark.asyncio
@patch("app.loader.glb.DEVICE", glb.DEVICE) # Mock the device to 'cpu'
async def test_zeroshot():
request_data = {
"input": "Test input",
@ -63,6 +75,7 @@ async def test_zeroshot():
# Unit test for the hallucination endpoint
@pytest.mark.asyncio
@patch("app.loader.glb.DEVICE", glb.DEVICE) # Mock the device to 'cpu'
async def test_hallucination():
request_data = {
"prompt": "Test hallucination",
@ -79,6 +92,7 @@ async def test_hallucination():
# Unit test for the chat completion endpoint
@pytest.mark.asyncio
@patch("app.loader.glb.DEVICE", glb.DEVICE) # Mock the device to 'cpu'
async def test_chat_completion():
async with httpx.AsyncClient(app=app, base_url="http://test") as client:
request_data = {