refactor tests for api

Author: isabelizimm

vetiver/handlers/base.py

@@ -1,3 +1,4 @@
+from typing import List
 from functools import singledispatch
 from contextlib import suppress
 import pandas as pd
@@ -119,8 +120,7 @@ def handler_startup():
         """
         ...
 
-    @singledispatch
-    def _prepare_data(self, pred_data) -> pd.DataFrame:
+    def _process_input(self, pred_data) -> pd.DataFrame:
         """Convert prototype to dataframe data
 
         Parameters
@@ -133,19 +133,10 @@ def _prepare_data(self, pred_data) -> pd.DataFrame:
         pd.DataFrame
             BaseModel data translated into DataFrame
         """
-        served_data = []
-        for key, value in pred_data:
-            served_data.append(value)
 
-        return served_data
+        new = _prepare_data(pred_data)
 
-    @_prepare_data.register
-    def _basemodel(self, pred_data: pydantic.BaseModel):
-        return pd.DataFrame([dict(s) for s in pred_data])
-
-    @_prepare_data.register
-    def _list(self, pred_data: list):
-        return pd.DataFrame([dict(s) for s in pred_data])
+        return new
 
     def handler_predict(self, input_data, check_prototype):
         """Generates method for /predict endpoint in VetiverAPI
@@ -179,3 +170,37 @@ def _(model: BaseHandler, prototype_data):
         model.prototype_data = prototype_data
 
     return model
+
+
+@singledispatch
+def _prepare_data(pred_data) -> pd.DataFrame:
+    """Convert prototype to dataframe data
+
+    Parameters
+    ----------
+    pred_data : pydantic.BaseModel
+        User data from given to API endpoint
+
+    Returns
+    -------
+    pd.DataFrame
+        BaseModel data translated into DataFrame
+    """
+
+    raise TypeError("Data should be list, dict, pd.DataFrame")
+
+
+@_prepare_data.register(pydantic.BaseModel)
+@_prepare_data.register(List)
+def _basemodel_list_data(pred_data):
+
+    return pd.DataFrame([dict(s) for s in pred_data])
+
+
+# @_prepare_data.register(dict)
+# def _dict_data(pred_data):
+#     served_data = []
+#     for key, value in pred_data:
+#         served_data.append(value)
+
+#     return served_data

vetiver/handlers/sklearn.py

@@ -34,7 +34,7 @@ def handler_predict(self, input_data, check_prototype):
             Prediction from model
         """
         if check_prototype:
-            input_data = self._prepare_data(input_data)
+            input_data = self._process_input(input_data)
 
         if not check_prototype or isinstance(input_data, pd.DataFrame):
             prediction = self.model.predict(input_data)

vetiver/handlers/statsmodels.py

@@ -41,9 +41,9 @@ def handler_predict(self, input_data, check_prototype):
         """
         if not sm_exists:
             raise ImportError("Cannot import `statsmodels`")
-        
+
         if check_prototype:
-            input_data = self._prepare_data(input_data)
+            input_data = self._process_input(input_data)
         if isinstance(input_data, (list, pd.DataFrame)):
             prediction = self.model.predict(input_data)
         else:

vetiver/handlers/torch.py

@@ -41,9 +41,9 @@ def handler_predict(self, input_data, check_prototype):
         """
         if not torch_exists:
             raise ImportError("Cannot import `torch`.")
-        
+
         if check_prototype:
-            input_data = self._prepare_data(input_data)
+            input_data = self._process_input(input_data)
             input_data = np.array(input_data, dtype=np.array(self.prototype_data).dtype)
             prediction = self.model(torch.from_numpy(input_data))
 

vetiver/handlers/xgboost.py

@@ -44,8 +44,8 @@ def handler_predict(self, input_data, check_prototype):
             raise ImportError("Cannot import `xgboost`")
 
         if check_prototype:
-            input_data = self._prepare_data(input_data)
-            
+            input_data = self._process_input(input_data)
+
         if not isinstance(input_data, pd.DataFrame):
             try:
                 input_data = pd.DataFrame(input_data)

vetiver/server.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, List, Union
+from typing import Callable, List, Union
 from urllib.parse import urljoin
 
 import httpx
@@ -167,15 +167,7 @@ def vetiver_post(self, endpoint_fx: Callable, endpoint_name: str = None, **kw):
         if self.check_prototype is True:
 
             @self.app.post(urljoin("/", endpoint_name), name=endpoint_name)
-            async def custom_endpoint(
-                input_data: Union[self.model.prototype, List[self.model.prototype]]
-            ):
-
-                # if isinstance(input_data, List):
-                #     served_data = _batch_data(input_data)
-                # else:
-                #     served_data = _prepare_data(input_data)
-
+            async def custom_endpoint(input_data: List[self.model.prototype]):
                 new = endpoint_fx(input_data, **kw)
                 return {endpoint_name: new.tolist()}
 
@@ -261,13 +253,13 @@ def predict(endpoint, data: Union[dict, pd.DataFrame, pd.Series], **kw) -> pd.Da
     # TO DO: dispatch
 
     if isinstance(data, pd.DataFrame):
-        data_json = data.to_json(orient="records")
-        response = requester.post(endpoint, data=data_json, **kw)
+        response = requester.post(endpoint, data=data.to_json(orient="records"), **kw)
     elif isinstance(data, pd.Series):
-        data_dict = data.to_json()
-        response = requester.post(endpoint, data=data_dict, **kw)
-    elif isinstance(data, dict):
+        response = requester.post(endpoint, data=data.to_json(), **kw)
+    elif isinstance(data, list):
         response = requester.post(endpoint, json=data, **kw)
+    elif isinstance(data, dict):
+        response = requester.post(endpoint, json=[data], **kw)
     else:
         response = requester.post(endpoint, json=data, **kw)