add tests

bananaml · Oct 25, 2023 · 7590044 · 7590044
1 parent 0580ba9
commit 7590044
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -7,7 +7,7 @@
 setup(
     name='potassium',
     packages=['potassium'],
-    version='0.3.1',
+    version='0.3.2',
     license='Apache License 2.0',
     # Give a short description about your library
     description='The potassium package is a flask-like HTTP server for serving large AI models',

diff --git a/tests/test_endpoints.py b/tests/test_endpoints.py
@@ -208,3 +208,9 @@ def handler(context: dict, request: potassium.Request) -> potassium.Response:
     res = client.post("/_k/warmup", json={})
     assert res.status_code == 200
     assert res.json == {"warm": True}
+
+    res = client.get("/__status__", json={})
+    assert res.status_code == 200
+    assert res.json is not None
+    assert res.json["gpu_available"] == True
+    assert res.json["sequence_number"] == 1