def test_all_y(mocker, raw_model, model_id, X, y, fit_params, predict_params, needs_proba):
model = TravaModel(raw_model=raw_model, model_id=model_id)
predict_proba_train = mocker.Mock()
if needs_proba:
raw_model.predict_proba.return_value = predict_proba_train
y_pred_train = mocker.Mock()
raw_model.predict.return_value = y_pred_train
model.fit(X=X, y=y, fit_params=fit_params, predict_params=predict_params)
predict_proba_test = mocker.Mock()
if needs_proba:
raw_model.predict_proba.return_value = predict_proba_test
y_pred_test = mocker.Mock()
raw_model.predict.return_value = y_pred_test
X_test = mocker.Mock()
y_test = mocker.Mock()
model.predict(X=X_test, y=y_test)
assert model.y_pred(for_train=True) == y_pred_train
assert model.y_pred(for_train=False) == y_pred_test
assert model.y(for_train=True) == y
assert model.y(for_train=False) == y_test
if needs_proba:
assert model.y_pred_proba(for_train=True) == predict_proba_train
assert model.y_pred_proba(for_train=False) == predict_proba_test
def test_copy(mocker, model_id, use_existing_model, only_fit):
raw_model = mocker.Mock()
model = TravaModel(raw_model=raw_model, model_id=model_id)
y_train = np.array([0, 0, 1])
y_train_pred = np.array([1, 2, 3])
y_train_pred_proba = np.array([3, 4, 5])
y_test = np.array([-2, 3, 5])
y_test_pred = np.array([6, 7, 8])
y_test_pred_proba = np.array([9, 10, 11])
fit_params = {"1": 2}
predict_params = {"2": 3}
fit_time = 123
predict_time = 434
model._y_train = y_train
model._y_train_pred = y_train_pred
model._y_train_pred_proba = y_train_pred_proba
model._y_test = y_test
model._y_test_pred = y_test_pred
model._y_test_pred_proba = y_test_pred_proba
model._fit_params = fit_params
model._predict_params = predict_params
model._fit_time = fit_time
model._predict_time = predict_time
model_copy_id = model_id + "_copy"
existing_model = None
existing_model_id = "existing_model"
# what a mess... but should work
if use_existing_model:
existing_model = TravaModel(raw_model=raw_model, model_id=existing_model_id)
model_copy = model.copy(existing_model=existing_model, only_fit=only_fit)
else:
model_copy = model.copy(model_id=model_copy_id, only_fit=only_fit)
if use_existing_model:
assert model_copy.model_id == existing_model_id
else:
assert model_copy.model_id == model_copy_id
y = model.y(for_train=True)
copy_y = model_copy.y(for_train=True)
assert np.array_equal(y, copy_y)
assert model.fit_params == model_copy.fit_params
assert model.fit_time == model_copy.fit_time
if use_existing_model:
assert existing_model == model_copy
else:
assert existing_model != model_copy
if only_fit:
assert model_copy.y(for_train=False) is None
assert model_copy.y_pred(for_train=False) is None
assert model_copy.y_pred_proba(for_train=False) is None
assert model_copy.predict_time is None
assert model_copy.predict_params == {}
else:
assert np.array_equal(model.y(for_train=False), model_copy.y(for_train=False))
assert np.array_equal(model.y_pred(for_train=False), model_copy.y_pred(for_train=False))
assert np.array_equal(model.y_pred_proba(for_train=False), model_copy.y_pred_proba(for_train=False))
assert model.predict_time == model_copy.predict_time
assert model.predict_params == model_copy.predict_params