KBaba7/llama.cpp

1import pytest2from openai import OpenAI3from utils import *4 5server = ServerPreset.tinyllama2()6 7TEST_API_KEY = "sk-this-is-the-secret-key"8 9@pytest.fixture(scope="module", autouse=True)10def create_server():11    global server12    server = ServerPreset.tinyllama2()13    server.api_key = TEST_API_KEY14 15 16@pytest.mark.parametrize("endpoint", ["/health", "/models"])17def test_access_public_endpoint(endpoint: str):18    global server19    server.start()20    res = server.make_request("GET", endpoint)21    assert res.status_code == 20022    assert "error" not in res.body23 24 25@pytest.mark.parametrize("api_key", [None, "invalid-key"])26def test_incorrect_api_key(api_key: str):27    global server28    server.start()29    res = server.make_request("POST", "/completions", data={30        "prompt": "I believe the meaning of life is",31    }, headers={32        "Authorization": f"Bearer {api_key}" if api_key else None,33    })34    assert res.status_code == 40135    assert "error" in res.body36    assert res.body["error"]["type"] == "authentication_error"37 38 39def test_correct_api_key():40    global server41    server.start()42    res = server.make_request("POST", "/completions", data={43        "prompt": "I believe the meaning of life is",44    }, headers={45        "Authorization": f"Bearer {TEST_API_KEY}",46    })47    assert res.status_code == 20048    assert "error" not in res.body49    assert "content" in res.body50 51 52def test_openai_library_correct_api_key():53    global server54    server.start()55    client = OpenAI(api_key=TEST_API_KEY, base_url=f"http://{server.server_host}:{server.server_port}")56    res = client.chat.completions.create(57        model="gpt-3.5-turbo",58        messages=[59            {"role": "system", "content": "You are a chatbot."},60            {"role": "user", "content": "What is the meaning of life?"},61        ],62    )63    assert len(res.choices) == 164 65 66@pytest.mark.parametrize("origin,cors_header,cors_header_value", [67    ("localhost", "Access-Control-Allow-Origin", "localhost"),68    ("web.mydomain.fr", "Access-Control-Allow-Origin", "web.mydomain.fr"),69    ("origin", "Access-Control-Allow-Credentials", "true"),70    ("web.mydomain.fr", "Access-Control-Allow-Methods", "GET, POST"),71    ("web.mydomain.fr", "Access-Control-Allow-Headers", "*"),72])73def test_cors_options(origin: str, cors_header: str, cors_header_value: str):74    global server75    server.start()76    res = server.make_request("OPTIONS", "/completions", headers={77        "Origin": origin,78        "Access-Control-Request-Method": "POST",79        "Access-Control-Request-Headers": "Authorization",80    })81    assert res.status_code == 20082    assert cors_header in res.headers83    assert res.headers[cors_header] == cors_header_value84