diff --git a/tests/integration/projects/general/service.py b/tests/integration/projects/general/service.py index b6c23c28d5b..d7565b4bf94 100644 --- a/tests/integration/projects/general/service.py +++ b/tests/integration/projects/general/service.py @@ -70,23 +70,20 @@ def predict_json(self, input_datas): def customezed_route(self, input_datas): return input_datas - @bentoml.api( - input=JsonInput( - request_schema={ - "application/json": { - "schema": { - "type": "object", - "required": ["field1", "field2"], - "properties": { - "field1": {"type": "string"}, - "field2": {"type": "uuid"}, - }, - }, - } - } - ), - batch=True, - ) + CUSTOM_SCHEMA = { + "application/json": { + "schema": { + "type": "object", + "required": ["field1", "field2"], + "properties": { + "field1": {"type": "string"}, + "field2": {"type": "uuid"}, + }, + }, + } + } + + @bentoml.api(input=JsonInput(request_schema=CUSTOM_SCHEMA), batch=True) def customezed_schema(self, input_datas): return input_datas diff --git a/tests/integration/projects/general_non_batch/tests/test_meta.py b/tests/integration/projects/general_non_batch/tests/test_meta.py index f394706c429..f2d6522f8ab 100644 --- a/tests/integration/projects/general_non_batch/tests/test_meta.py +++ b/tests/integration/projects/general_non_batch/tests/test_meta.py @@ -11,9 +11,13 @@ async def test_api_server_meta(host): @pytest.mark.asyncio async def test_customized_request_schema(host): + def has_customized_schema(doc_bytes): + json_str = doc_bytes.decode() + return "field1" in json_str + await pytest.assert_request( "GET", f"http://{host}/docs.json", headers=(("Content-Type", "application/json"),), - assert_data="", + assert_data=has_customized_schema, )