test serialization

Author: isabelizimm

setup.cfg

@@ -48,3 +48,6 @@ dev =
 
 torch =
     torch
+
+statsmodels =
+    statsmodels

vetiver/handlers/statsmodels.py

@@ -19,9 +19,7 @@ class StatsmodelsHandler(BaseHandler):
         a trained sklearn model
     """
 
-    model_class = staticmethod(
-        lambda: statsmodels.regression.linear_model.RegressionResultsWrapper
-    )
+    model_class = staticmethod(lambda: statsmodels.base.wrapper.ResultsWrapper)
 
     def __init__(self, model, ptype_data):
         super().__init__(model, ptype_data)

vetiver/tests/test_statsmodels.py

@@ -1,128 +1,43 @@
-# import pytest
+import pytest
 
-# sm = pytest.importorskip("statsmodels", reason="statsmodels library not installed")
+sm = pytest.importorskip("statsmodels.api", reason="statsmodels library not installed")
 
-# import numpy as np  # noqa
-# from fastapi.testclient import TestClient  # noqa
+statsmodels = pytest.importorskip(
+    "statsmodels", reason="statsmodels library not installed"
+)
 
-# from vetiver.vetiver_model import VetiverModel  # noqa
-# from vetiver import VetiverAPI  # noqa
+import numpy as np  # noqa
+from fastapi.testclient import TestClient  # noqa
 
+import vetiver  # noqa
 
-# def _build_sm():
 
-#     input_size = 1
-#     output_size = 1
+@pytest.fixture
+def build_sm():
 
-#     x_train = np.array(
-#         [
-#             [3.3],
-#             [4.4],
-#             [5.5],
-#             [6.71],
-#             [6.93],
-#             [4.168],
-#             [9.779],
-#             [6.182],
-#             [7.59],
-#             [2.167],
-#             [7.042],
-#             [10.791],
-#             [5.313],
-#             [7.997],
-#             [3.1],
-#         ],
-#         dtype=np.float32,
-#     )
+    X, y = vetiver.get_mock_data()
+    glm = sm.GLM(y, X).fit()
 
-#     torch_model = sm.nn.Linear(input_size, output_size)
-#     return x_train, torch_model
+    v = vetiver.VetiverModel(glm, "glm", X)
+    return v
 
 
-# def test_vetiver_build():
+def test_vetiver_build(build_sm):
+    api = vetiver.VetiverAPI(build_sm)
+    client = TestClient(api.app)
+    data = {"B": 0, "C": 0, "D": 0}
+    response = client.post("/predict", json=data)
+    assert response.status_code == 200, response.text
+    assert response.json() == {"prediction": [0.0]}, response.json()
 
-#     x_train, torch_model = _build_sm()
 
-#     vt2 = VetiverModel(
-#         model=torch_model,
-#         ptype_data=x_train,
-#         model_name="torch",
-#         versioned=None,
-#         description=None,
-#         metadata=None,
-#     )
+def test_serialize(build_sm):
+    import pins
 
-#     assert vt2.model == torch_model
-
-
-# def test_sm_predict_ptype():
-#     torch.manual_seed(3)
-#     x_train, torch_model = _build_sm()
-#     v = VetiverModel(torch_model, model_name="torch", ptype_data=x_train)
-#     v_api = VetiverAPI(v)
-
-#     client = TestClient(v_api.app)
-#     data = {"0": 3.3}
-#     response = client.post("/predict", json=data)
-
-#     assert response.status_code == 200, response.text
-#     assert response.json() == {"prediction": [-4.060722351074219]}, response.text
-
-
-# def test_sm_predict_ptype_batch():
-
-#     x_train, torch_model = _build_sm()
-#     v = VetiverModel(torch_model, model_name="torch", ptype_data=x_train)
-#     v_api = VetiverAPI(v)
-
-#     client = TestClient(v_api.app)
-#     data = [{"0": 3.3}, {"0": 3.3}]
-#     response = client.post("/predict", json=data)
-
-#     assert response.status_code == 200, response.text
-#     assert response.json() == {
-#         "prediction": [[-4.060722351074219], [-4.060722351074219]]
-#     }, response.text
-
-
-# def test_sm_predict_ptype_error():
-
-#     x_train, torch_model = _build_sm()
-#     v = VetiverModel(torch_model, model_name="torch", ptype_data=x_train)
-#     v_api = VetiverAPI(v)
-
-#     client = TestClient(v_api.app)
-#     data = {"0": "bad"}
-#     response = client.post("/predict", json=data)
-
-#     assert response.status_code == 422, response.text  # value is not a valid float
-
-
-# def test_sm_predict_no_ptype_batch():
-
-#     x_train, torch_model = _build_sm()
-#     v = VetiverModel(torch_model, model_name="torch")
-#     v_api = VetiverAPI(v, check_ptype=False)
-
-#     client = TestClient(v_api.app)
-#     data = [[3.3], [3.3]]
-#     response = client.post("/predict", json=data)
-#     assert response.status_code == 200, response.text
-#     assert response.json() == {
-#         "prediction": [[-4.060722351074219], [-4.060722351074219]]
-#     }, response.text
-
-
-# def test_sm_predict_no_ptype():
-
-#     x_train, torch_model = _build_sm()
-#     v = VetiverModel(torch_model, model_name="torch")
-#     v_api = VetiverAPI(v, check_ptype=False)
-
-#     client = TestClient(v_api.app)
-#     data = [[3.3]]
-#     response = client.post("/predict", json=data)
-#     assert response.status_code == 200, response.text
-#     assert response.json() == {"prediction": [[-4.060722351074219]]}, response.text
-
-# def test_pin_sm():
+    board = pins.board_temp(allow_pickle_read=True)
+    vetiver.vetiver_pin_write(board=board, model=build_sm)
+    assert isinstance(
+        board.pin_read("glm"),
+        statsmodels.genmod.generalized_linear_model.GLMResultsWrapper,
+    )
+    board.pin_delete("glm")