def test_get_model(mocker, raw_model, X, y, needs_proba):
model = TravaModel(raw_model=raw_model, model_id=model_id)
assert model.get_model(for_train=True) == raw_model
assert model.get_model(for_train=False) == raw_model
y_predict_proba = mocker.Mock()
if needs_proba:
raw_model.predict_proba.return_value = y_predict_proba
y_pred = mocker.Mock()
raw_model.predict.return_value = y_pred
model.fit(X=X, y=y)
model.predict(X=X, y=y)
model.unload_model()
train_cached_model = model.get_model(for_train=True)
test_cached_model = model.get_model(for_train=False)
assert train_cached_model != raw_model
assert test_cached_model != raw_model
assert train_cached_model.predict(X) == y_pred
if needs_proba:
assert train_cached_model.predict_proba(X) == y_predict_proba
def test_get_model_unload(mocker, raw_model, for_train):
trava_model = TravaModel(raw_model=raw_model, model_id=model_id)
trava_model.unload_model()
with pytest.raises(ValueError):
trava_model.get_model(for_train=for_train)
if for_train:
y_pred_key = "_y_train_pred"
else:
y_pred_key = "_y_test_pred"
y_pred_mock = mocker.Mock()
mocker.patch.object(trava_model, y_pred_key, y_pred_mock)
assert trava_model.get_model(for_train=for_train).predict(X=None) == y_pred_mock
def test_unload(mocker, model_id):
raw_model = mocker.Mock()
model = TravaModel(raw_model=raw_model, model_id=model_id)
model.unload_model()
assert not model.raw_model