add test

ngxson · ngxson · commit 7747225cedb7 · 2024-12-10T17:06:10.000+01:00
diff --git a/examples/server/tests/unit/test_basic.py b/examples/server/tests/unit/test_basic.py
@@ -1,4 +1,5 @@
 import pytest
+import requests
 from utils import *
 
 server = ServerPreset.tinyllama2()
@@ -76,3 +77,19 @@ def test_load_split_model():
     })
     assert res.status_code == 200
     assert match_regex("(little|girl)+", res.body["content"])
+
+
+def test_no_webui():
+    global server
+    # default: webui enabled
+    server.start()
+    url = f"http://{server.server_host}:{server.server_port}"
+    res = requests.get(url)
+    assert res.status_code == 200
+    assert "<html>" in res.text
+
+    # with --no-webui
+    server.no_webui = True
+    server.start()
+    res = requests.get(url)
+    assert res.status_code == 404
diff --git a/examples/server/tests/utils.py b/examples/server/tests/utils.py
@@ -72,6 +72,7 @@ class ServerProcess:
     disable_ctx_shift: int | None = False
     draft_min: int | None = None
     draft_max: int | None = None
+    no_webui: bool | None = None
 
     # session variables
     process: subprocess.Popen | None = None
@@ -158,6 +159,8 @@ def start(self, timeout_seconds: int = 10) -> None:
             server_args.extend(["--draft-max", self.draft_max])
         if self.draft_min:
             server_args.extend(["--draft-min", self.draft_min])
+        if self.no_webui:
+            server_args.append("--no-webui")
 
         args = [str(arg) for arg in [server_path, *server_args]]
         print(f"bench: starting server with: {' '.join(args)}")
