no ptype and batch tests

Author: isabelizimm

vetiver/handlers/statsmodels.py

@@ -58,9 +58,8 @@ def handler_predict(self, input_data, check_ptype):
         prediction
             Prediction from model
         """
-        if not check_ptype:
-            input_data = pd.DataFrame(input_data)
-        if isinstance(input_data, pd.DataFrame):
+
+        if isinstance(input_data, (list, pd.DataFrame)):
             prediction = self.model.predict(input_data)
         else:
             prediction = self.model.predict([input_data])

vetiver/server.py

@@ -83,12 +83,8 @@ async def prediction(
             @app.post("/predict")
             async def prediction(input_data: Request):
                 y = await input_data.json()
-                from io import BytesIO  # noqa
 
-                df = pd.read_csv(BytesIO(y))
-                prediction = self.model.handler_predict(
-                    df, check_ptype=self.check_ptype
-                )
+                prediction = self.model.handler_predict(y, check_ptype=self.check_ptype)
 
                 return {"prediction": prediction.tolist()}
 

vetiver/tests/test_statsmodels.py

@@ -7,6 +7,7 @@
 )
 
 import numpy as np  # noqa
+import pandas as pd  # noqa
 from fastapi.testclient import TestClient  # noqa
 
 import vetiver  # noqa
@@ -25,20 +26,28 @@ def build_sm():
 def test_vetiver_build(build_sm):
     api = vetiver.VetiverAPI(build_sm)
     client = TestClient(api.app)
-    data = {"B": 0, "C": 0, "D": 0}
+    data = [{"B": 0, "C": 0, "D": 0}]
 
     response = vetiver.predict(endpoint=client, data=data)
 
     assert response.iloc[0, 0] == 0.0
     assert len(response) == 1
 
 
+def test_batch(build_sm):
+    api = vetiver.VetiverAPI(build_sm)
+    client = TestClient(api.app)
+    data = pd.DataFrame(np.random.randint(0, 100, size=(100, 4)), columns=list("ABCD"))
+
+    response = vetiver.predict(endpoint=client, data=data)
+
+    assert len(response) == 100
+
+
 def test_no_ptype(build_sm):
-    v = build_sm
-    v.ptype = None
-    api = vetiver.VetiverAPI(v)
+    api = vetiver.VetiverAPI(build_sm, check_ptype=False)
     client = TestClient(api.app)
-    data = 0, 0, 0
+    data = [0, 0, 0]
 
     response = vetiver.predict(endpoint=client, data=data)