fix tests

Lightning-AI · Dec 11, 2024 · b47bc10 · b47bc10
1 parent 863fbcf
commit b47bc10
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 10 deletions.
diff --git a/tests/test_litapi.py b/tests/test_litapi.py
@@ -175,7 +175,7 @@ def predict():
 
 
 def test_encode_response_with_custom_spec_api():
-    class CustomSpecAPI(ls.test_examples.TestAPI):
+    class CustomSpecAPI(ls.OpenAISpec):
         def encode_response(self, output_stream):
             for output in output_stream:
                 yield {"content": output}

diff --git a/tests/test_specs.py b/tests/test_specs.py
@@ -176,15 +176,11 @@ def encode_response(self, output):
 
 @pytest.mark.asyncio
 async def test_openai_spec_validation(openai_request_data):
-    server = ls.LitServer(IncorrectAPI1(), spec=OpenAISpec())
-    with pytest.raises(ValueError, match="predict is not a generator"), wrap_litserve_start(server) as server:
-        async with LifespanManager(server.app) as manager:
-            await manager.shutdown()
-
-    server = ls.LitServer(IncorrectAPI2(), spec=OpenAISpec())
-    with pytest.raises(ValueError, match="encode_response is not a generator"), wrap_litserve_start(server) as server:
-        async with LifespanManager(server.app) as manager:
-            await manager.shutdown()
+    with pytest.raises(ValueError, match="predict is not a generator"):
+        ls.LitServer(IncorrectAPI1(), spec=OpenAISpec())
+
+    with pytest.raises(ValueError, match="encode_response is not a generator"):
+        ls.LitServer(IncorrectAPI2(), spec=OpenAISpec())
 
 
 class PrePopulatedAPI(ls.LitAPI):