Felipe97/llama-cpp-compiled
01.2k
1import pytest2from openai import OpenAI3from utils import *4import threading5from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer6 7server = ServerPreset.tinyllama2()8 9TEST_API_KEY = "sk-this-is-the-secret-key"10 11@pytest.fixture(autouse=True)12def create_server():13 global server14 server = ServerPreset.tinyllama2()15 server.api_key = TEST_API_KEY16 17 18@pytest.mark.parametrize("endpoint", ["/health"])19def test_access_public_endpoint(endpoint: str):20 global server21 server.start()22 res = server.make_request("GET", endpoint)23 assert res.status_code == 20024 assert "error" not in res.body25 26 27def test_access_static_assets_without_api_key():28 """Static web UI assets should not require API key authentication (issue #21229)"""29 global server30 server.start()31 for path in ["/", "/sw.js", "/manifest.webmanifest", "/_app/version.json"]:32 res = server.make_request("GET", path)33 assert res.status_code == 200, f"Expected 200 for {path}, got {res.status_code}"34 35 36@pytest.mark.parametrize("api_key", [None, "invalid-key"])37def test_incorrect_api_key(api_key: str):38 global server39 server.start()40 res = server.make_request("POST", "/completions", data={41 "prompt": "I believe the meaning of life is",42 }, headers={43 "Authorization": f"Bearer {api_key}" if api_key else None,44 })45 assert res.status_code == 40146 assert "error" in res.body47 assert res.body["error"]["type"] == "authentication_error"48 49 50def test_correct_api_key():51 global server52 server.start()53 res = server.make_request("POST", "/completions", data={54 "prompt": "I believe the meaning of life is",55 }, headers={56 "Authorization": f"Bearer {TEST_API_KEY}",57 })58 assert res.status_code == 20059 assert "error" not in res.body60 assert "content" in res.body61 62 63def test_correct_api_key_anthropic_header():64 global server65 server.start()66 res = server.make_request("POST", "/completions", data={67 "prompt": "I believe the meaning of life is",68 }, headers={69 "X-Api-Key": TEST_API_KEY,70 })71 assert res.status_code == 20072 assert "error" not in res.body73 assert "content" in res.body74 75 76def test_openai_library_correct_api_key():77 global server78 server.start()79 client = OpenAI(api_key=TEST_API_KEY, base_url=f"http://{server.server_host}:{server.server_port}")80 res = client.chat.completions.create(81 model="gpt-3.5-turbo",82 messages=[83 {"role": "system", "content": "You are a chatbot."},84 {"role": "user", "content": "What is the meaning of life?"},85 ],86 )87 assert len(res.choices) == 188 89 90@pytest.mark.parametrize("origin,cors_header,cors_header_value", [91 ("localhost", "Access-Control-Allow-Origin", "localhost"),92 ("web.mydomain.fr", "Access-Control-Allow-Origin", "web.mydomain.fr"),93 ("origin", "Access-Control-Allow-Credentials", "true"),94 ("web.mydomain.fr", "Access-Control-Allow-Methods", "GET, POST, DELETE, OPTIONS"),95 ("web.mydomain.fr", "Access-Control-Allow-Headers", "*"),96])97def test_cors_options(origin: str, cors_header: str, cors_header_value: str):98 global server99 server.start()100 res = server.make_request("OPTIONS", "/completions", headers={101 "Origin": origin,102 "Access-Control-Request-Method": "POST",103 "Access-Control-Request-Headers": "Authorization",104 })105 assert res.status_code == 200106 assert cors_header in res.headers107 assert res.headers[cors_header] == cors_header_value108 109 110@pytest.mark.parametrize("origin", [111 "http://localhost",112 "http://localhost:8080",113 "http://127.0.0.1",114 "http://127.0.0.1:3000",115 "http://[::1]",116 "http://[::1]:3000",117])118def test_cors_origins_localhost_reflects(origin: str):119 global server120 server = ServerPreset.router()121 server.cors_origins = "localhost"122 server.start()123 res = server.make_request("OPTIONS", "/completions", headers={124 "Origin": origin,125 "Access-Control-Request-Method": "POST",126 "Access-Control-Request-Headers": "Authorization",127 })128 assert res.status_code == 200129 assert res.headers["Access-Control-Allow-Origin"] == origin130 131 132@pytest.mark.parametrize("origin", [133 "http://web.mydomain.fr",134 "http://evil.com",135 "http://notlocalhost",136 "http://localhost.evil.com",137])138def test_cors_origins_localhost_rejects(origin: str):139 global server140 server = ServerPreset.router()141 server.cors_origins = "localhost"142 server.start()143 res = server.make_request("OPTIONS", "/completions", headers={144 "Origin": origin,145 "Access-Control-Request-Method": "POST",146 "Access-Control-Request-Headers": "Authorization",147 })148 assert res.status_code == 200149 assert "Access-Control-Allow-Origin" not in res.headers150 151 152def test_cors_origins_defaults_to_localhost_with_tools_enabled():153 global server154 server = ServerPreset.router()155 server.server_tools = "all"156 server.start()157 res = server.make_request("OPTIONS", "/completions", headers={158 "Origin": "http://localhost:8080",159 "Access-Control-Request-Method": "POST",160 "Access-Control-Request-Headers": "Authorization",161 })162 assert res.status_code == 200163 assert res.headers["Access-Control-Allow-Origin"] == "http://localhost:8080"164 165 res = server.make_request("OPTIONS", "/completions", headers={166 "Origin": "http://evil.com",167 "Access-Control-Request-Method": "POST",168 "Access-Control-Request-Headers": "Authorization",169 })170 assert res.status_code == 200171 assert "Access-Control-Allow-Origin" not in res.headers172 173 174def test_cors_proxy_only_forwards_explicit_proxy_headers():175 class CaptureHeadersHandler(BaseHTTPRequestHandler):176 def do_GET(self):177 self.server.captured_headers = dict(self.headers)178 self.send_response(200)179 self.end_headers()180 self.wfile.write(b"ok")181 182 def log_message(self, format, *args):183 pass184 185 target = ThreadingHTTPServer(("127.0.0.1", 0), CaptureHeadersHandler)186 target.captured_headers = {}187 target_thread = threading.Thread(target=target.serve_forever, daemon=True)188 target_thread.start()189 190 try:191 server = ServerPreset.tinyllama2()192 server.api_key = TEST_API_KEY193 server.ui_mcp_proxy = True194 server.start()195 196 res = server.make_request("GET", f"/cors-proxy?url=http://127.0.0.1:{target.server_port}/capture", headers={197 "Authorization": f"Bearer {TEST_API_KEY}",198 "Proxy-Authorization": "Basic secret",199 "X-Api-Key": TEST_API_KEY,200 "Cookie": "session=secret",201 "x-llama-server-proxy-header-accept": "application/json",202 "x-llama-server-proxy-header-authorization": "Bearer explicit",203 })204 205 assert res.status_code == 200206 captured = {key.lower(): value for key, value in target.captured_headers.items()}207 assert captured["accept"] == "application/json"208 assert captured["authorization"] == "Bearer explicit"209 assert "proxy-authorization" not in captured210 assert "x-api-key" not in captured211 assert "cookie" not in captured212 finally:213 target.shutdown()214 target.server_close()215 216 217@pytest.mark.parametrize(218 "media_path, image_url, success",219 [220 (None, "file://mtmd/test-1.jpeg", False), # disabled media path, should fail221 ("../../../tools", "file://mtmd/test-1.jpeg", True),222 ("../../../tools", "file:////mtmd//test-1.jpeg", True), # should be the same file as above223 ("../../../tools", "file://mtmd/notfound.jpeg", False), # non-existent file224 ("../../../tools", "file://../mtmd/test-1.jpeg", False), # no directory traversal225 ]226)227def test_local_media_file(media_path, image_url, success,):228 server = ServerPreset.tinygemma3()229 server.media_path = media_path230 server.start()231 res = server.make_request("POST", "/chat/completions", data={232 "max_tokens": 1,233 "messages": [234 {"role": "user", "content": [235 {"type": "text", "text": "test"},236 {"type": "image_url", "image_url": {237 "url": image_url,238 }},239 ]},240 ],241 })242 if success:243 assert res.status_code == 200244 else:245 assert res.status_code == 400246 