Added missing response_headers variable in predict call.

Signed-off-by: Andrews Arokiam <[email protected]>
vignesh-murugani2i · Nov 29, 2023 · d4fd668 · d4fd668
1 parent 000d79d
commit d4fd668
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/python/lgbserver/lgbserver/test_model.py b/python/lgbserver/lgbserver/test_model.py
@@ -56,10 +56,10 @@ def test_model():
         {'sepal_width_(cm)': 3.5}, {'petal_length_(cm)': 1.4},
         {'petal_width_(cm)': 0.2}, {'sepal_length_(cm)': 5.1}
     ]
-    response = model.predict({"inputs": [request, request]})
+    response, response_headers = model.predict({"inputs": [request, request]})
     assert numpy.argmax(response["predictions"][0]) == 2
 
-    response = model.predict({"instances": [request, request]})
+    response, response_headers = model.predict({"instances": [request, request]})
     assert numpy.argmax(response["predictions"][0]) == 2
 
     # test v2 handler

diff --git a/python/test_resources/graph/success_200_isvc/model.py b/python/test_resources/graph/success_200_isvc/model.py
@@ -32,7 +32,8 @@ def load(self):
         self.ready = True
 
     def predict(self, payload: Union[Dict, InferRequest, ModelInferRequest], headers) -> Dict:
-        return {"message": "SUCCESS"}
+        response_headers = {}
+        return {"message": "SUCCESS"}, response_headers
 
 
 DEFAULT_MODEL_NAME = "model"