From b47bc10f7687da4d84f911f950199abf8bfc910d Mon Sep 17 00:00:00 2001 From: Aniket Maurya Date: Wed, 11 Dec 2024 22:13:58 +0000 Subject: [PATCH] fix tests --- tests/test_litapi.py | 2 +- tests/test_specs.py | 14 +++++--------- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/tests/test_litapi.py b/tests/test_litapi.py index 797ecbdf..6c92c018 100644 --- 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 index 360b376c..78b71980 100644 --- 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):