tests

Author: isabelizimm

vetiver/server.py

@@ -2,6 +2,7 @@
 from fastapi.responses import HTMLResponse, RedirectResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.openapi.utils import get_openapi
+from fastapi import testclient
 
 import uvicorn
 import requests
@@ -184,13 +185,19 @@ def predict(endpoint, data: Union[dict, pd.DataFrame, pd.Series], **kw):
     dict
         Key: endpoint_name Value: Output of endpoint_fx, in list format
     """
+    if isinstance(endpoint, testclient.TestClient):
+        requester = endpoint
+        endpoint = "/predict/"
+    else: 
+        requester = requests
+
     if isinstance(data, (pd.DataFrame,pd.Series)):
         data_json = data.to_json(orient="records")
-        response = requests.post(endpoint, data=data_json, **kw)
+        response = requester.post(endpoint, data=data_json, **kw)
     elif isinstance(data, dict):
-        response = requests.post(endpoint, json=data, **kw)
+        response = requester.post(endpoint, json=data, **kw)
     else:
-        raise TypeError(f"Accepted data types are dictionary or DataFrame, given type is {type(data)} \n {data}")
+        raise TypeError(f"Given type is {type(data)}")
 
     response_df = pd.DataFrame.from_dict(response.json())
 

vetiver/tests/test_predict.py

@@ -0,0 +1,115 @@
+import pytest
+
+import numpy as np
+import pandas as pd
+from fastapi.testclient import TestClient
+
+from vetiver import mock, VetiverModel, VetiverAPI
+from vetiver.server import predict
+
+
+def test_predict_sklearn_dict_ptype():
+    np.random.seed(500)
+    X, y = mock.get_mock_data()
+    model = mock.get_mock_model().fit(X, y)
+    v = VetiverModel(
+        model=model,
+        save_ptype=True,
+        ptype_data=X,
+        model_name="my_model",
+        versioned=None,
+        description="A regression model for testing purposes",
+    )
+    app = VetiverAPI(v, check_ptype=True)
+    client = TestClient(app.app)
+    data = {"B": 0, "C": 0, "D": 0}
+    
+    response = predict(endpoint=client, data=data)
+
+    assert isinstance(response, pd.DataFrame), response
+    assert response.iloc[0,0] == 44.47
+
+
+def test_predict_sklearn_no_ptype():
+    np.random.seed(500)
+    X, y = mock.get_mock_data()
+    model = mock.get_mock_model().fit(X, y)
+    v = VetiverModel(
+        model=model,
+        save_ptype=True,
+        ptype_data=X,
+        model_name="my_model",
+        versioned=None,
+        description="A regression model for testing purposes",
+    )
+    app = VetiverAPI(v, check_ptype=False)
+    client = TestClient(app.app)
+    
+    response = predict(endpoint=client, data=X)
+
+    assert isinstance(response, pd.DataFrame), response
+    assert response.iloc[0,0] == 44.47
+    assert len(response) == 100
+
+
+def test_predict_sklearn_df_check_ptype():
+    np.random.seed(500)
+    X, y = mock.get_mock_data()
+    model = mock.get_mock_model().fit(X, y)
+    v = VetiverModel(
+        model=model,
+        save_ptype=True,
+        ptype_data=X,
+        model_name="my_model",
+        versioned=None,
+        description="A regression model for testing purposes",
+    )
+    app = VetiverAPI(v, check_ptype=True)
+    client = TestClient(app.app)
+    
+    response = predict(endpoint=client, data=X)
+
+    assert isinstance(response, pd.DataFrame), response
+    assert response.iloc[0,0] == 44.47
+    assert len(response) == 100
+
+
+def test_predict_sklearn_df_no_ptype():
+    np.random.seed(500)
+    X, y = mock.get_mock_data()
+    model = mock.get_mock_model().fit(X, y)
+    v = VetiverModel(
+        model=model,
+        save_ptype=True,
+        ptype_data=X,
+        model_name="my_model",
+        versioned=None,
+        description="A regression model for testing purposes",
+    )
+    app = VetiverAPI(v, check_ptype=False)
+    client = TestClient(app.app)
+    
+    response = predict(endpoint=client, data=X)
+
+    assert isinstance(response, pd.DataFrame), response
+    assert response.iloc[0,0] == 44.47
+
+
+def test_predict_sklearn_type_error():
+    np.random.seed(500)
+    X, y = mock.get_mock_data()
+    model = mock.get_mock_model().fit(X, y)
+    v = VetiverModel(
+        model=model,
+        save_ptype=True,
+        ptype_data=X,
+        model_name="my_model",
+        versioned=None,
+        description="A regression model for testing purposes",
+    )
+    app = VetiverAPI(v, check_ptype=True)
+    client = TestClient(app.app)
+    data = '0,0,0'
+    
+    with pytest.raises(TypeError):
+        predict(endpoint=client, data=data)