tests: Add pickling tests for regression models
@@ -1480,6 +1480,14 @@ class TestRegressionFitters:
regression_models_sans_strata_model.append(CoxPHFitter(strata=["race", "paro", "mar", "wexp"]))
return regression_models_sans_strata_model
+ def test_pickle_serialization(self, rossi, regression_models):
+ for fitter in regression_models:
+ fitter.fit(rossi, "week", "arrest")
+
+ unpickled = pickle.loads(pickle.dumps(fitter))
+ dif = (fitter.durations - unpickled.durations).sum()
+ assert dif == 0
+
def test_dill_serialization(self, rossi, regression_models):
from dill import dumps, loads
@@ -1490,7 +1498,6 @@ class TestRegressionFitters:
dif = (fitter.durations - unpickled.durations).sum()
assert dif == 0
- @pytest.mark.xfail()
def test_joblib_serialization(self, rossi, regression_models):
from joblib import dump, load
@@ -1502,15 +1509,6 @@ class TestRegressionFitters:
dif = (fitter.durations - unpickled.durations).sum()
assert dif == 0
- @pytest.mark.xfail()
- def test_pickle(self, rossi, regression_models):
- from pickle import dump
-
- for fitter in regression_models:
- output = stringio()
- f = fitter.fit(rossi, "week", "arrest")
- dump(f, output)
-
def test_fit_will_accept_object_dtype_as_event_col(self, regression_models_sans_strata_model, rossi):
# issue #638
rossi["arrest"] = rossi["arrest"].astype(object)
