def test_save(self, tmp_path, monkeypatch):
coef_array = np.array([1, 1, 1, 1, 1])
intercept_array = np.array([2])
def mocktrain(self):
class MockModel:
@property
def coef_(self):
return coef_array
@property
def intercept_(self):
return intercept_array
self.model = MockModel()
monkeypatch.setattr(LinearRegression, "train", mocktrain)
model = LinearRegression(
tmp_path, experiment="one_month_forecast", normalize_y=False
)
model.train()
model.save_model()
assert (
tmp_path / "models/one_month_forecast/linear_regression/model.pkl"
).exists(), f"Model not saved!"
with (tmp_path / "models/one_month_forecast/linear_regression/model.pkl").open(
"rb"
) as f:
model_dict = pickle.load(f)
assert np.array_equal(
coef_array, model_dict["model"]["coef"]
), "Different coef array saved!"
assert np.array_equal(
intercept_array, model_dict["model"]["intercept"]
), "Different intercept array saved!"
assert (
model_dict["experiment"] == "one_month_forecast"
), "Different experiment saved!"
def test_save(self, tmp_path, monkeypatch):
coef_array = np.array([1, 1, 1, 1, 1])
intercept_array = np.array([2])
def mocktrain(self):
class MockModel:
@property
def coef_(self):
return coef_array
@property
def intercept_(self):
return intercept_array
self.model = MockModel()
monkeypatch.setattr(LinearRegression, 'train', mocktrain)
model = LinearRegression(tmp_path, experiment='one_month_forecast')
model.train()
model.save_model()
assert (tmp_path /
'models/one_month_forecast/linear_regression/model.pkl'
).exists(), f'Model not saved!'
with (tmp_path /
'models/one_month_forecast/linear_regression/model.pkl'
).open('rb') as f:
model_dict = pickle.load(f)
assert np.array_equal(coef_array, model_dict['model']['coef']), \
'Different coef array saved!'
assert np.array_equal(intercept_array, model_dict['model']['intercept']), \
'Different intercept array saved!'
assert model_dict[
'experiment'] == 'one_month_forecast', 'Different experiment saved!'