def test_modelpipeline_etl_init_err():
# If users don't have access to >= v2.0 temlates, then passing
# `etl` to a new ModelPipeline should produce a NotImplementedError.
mock_client = mock.MagicMock()
with pytest.raises(NotImplementedError):
_model.ModelPipeline(LogisticRegression(), 'test',
etl=LogisticRegression(),
client=mock_client)
def test_modelpipeline_init_newest():
mock_client = mock.MagicMock()
mock_client.templates.get_scripts.return_value = {}
etl = LogisticRegression()
mp = _model.ModelPipeline(LogisticRegression(),
'test',
etl=etl,
client=mock_client)
assert mp.etl == etl
# clean up
del _model._NEWEST_CIVISML_VERSION
def test_modelpipeline_init_err():
mock_client = mock.MagicMock()
r = Response({'content': None, 'status_code': 9999, 'reason': None})
mock_client.templates.get_scripts.side_effect = CivisAPIError(r)
with pytest.raises(NotImplementedError):
_model.ModelPipeline(LogisticRegression(),
'test',
etl=LogisticRegression(),
client=mock_client)
# clean up
del _model._NEWEST_CIVISML_VERSION
def test_modelpipeline_train_from_estimator(mock_ccr, mock_f2c):
# Provide a model as a pre-made model and make sure we can train.
mock_f2c.return_value = -21
est = LogisticRegression()
mp = _model.ModelPipeline(est, "dv")
mock1, mock2 = mock.Mock(), mock.Mock()
mock_ccr.return_value = 'res', mock1, mock2
assert 'res' == mp.train(file_id=7)
assert mp.train_result_ == 'res'
assert mock_f2c.call_count == 1 # Called once to store input Estimator
def test_modelpipeline_train_custom_etl(mock_ccr, mock_f2c, mp_setup):
# Provide a custom ETL estimator and make sure we can train.
mock_api = setup_client_mock()
etl = LogisticRegression()
mp = _model.ModelPipeline('wf', 'dv', client=mock_api, etl=etl)
mock_f2c.return_value = -21
mock1, mock2 = mock.Mock(), mock.Mock()
mock_ccr.return_value = 'res', mock1, mock2
assert 'res' == mp.train(file_id=7)
assert mp.train_result_ == 'res'
assert mock_f2c.call_count == 1 # Called once to store input Estimator
def test_modelpipeline_init_newest():
_model._CIVISML_TEMPLATE = None
mock_client = mock.MagicMock()
mock_client.templates.get_scripts.return_value = {}
etl = LogisticRegression()
mp = _model.ModelPipeline(LogisticRegression(),
'test',
etl=etl,
client=mock_client)
assert mp.etl == etl
assert mp.train_template_id == LATEST_TRAIN_TEMPLATE
# clean up
_model._CIVISML_TEMPLATE = None
def test_modelpipeline_train_custom_etl(mock_ccr, mock_f2c, mock_template_ids):
# Provide a custom ETL estimator and make sure we can train.
mock_api = create_client_mock_for_container_tests()
# training template ID 11111 >= 9968 for the etl arg to work
mock_template_ids.return_value = 11111, 22222, 33333
etl = LogisticRegression()
mp = _model.ModelPipeline('wf', 'dv', client=mock_api, etl=etl)
mock_f2c.return_value = -21
mock1, mock2 = mock.Mock(), mock.Mock()
mock_ccr.return_value = 'res', mock1, mock2
assert 'res' == mp.train(file_id=7)
assert mp.train_result_ == 'res'
assert mock_f2c.call_count == 1 # Called once to store input Estimator
def test_modelpipeline_etl_init_err():
# If users don't have access to >= v2.0 temlates, then passing
# `etl` to a new ModelPipeline should produce a NotImplementedError.
mock_client = mock.MagicMock()
r = Response({'content': None, 'status_code': 9999, 'reason': None})
def pre_2p0_template(id=None, **kwargs):
if id > 9113:
raise CivisAPIError(r)
return {}
mock_client.templates.get_scripts.side_effect = pre_2p0_template
with pytest.raises(NotImplementedError):
_model.ModelPipeline(LogisticRegression(),
'test',
etl=LogisticRegression(),
client=mock_client)
# clean up
_model._CIVISML_TEMPLATE = None
def test_modelpipeline_pickling_preserves_template_ids(
version, train_id, predict_id):
# Test that pickling a ModelPipeline object preserves the template IDs
# that have already been set during object instantiation.
with TemporaryDirectory() as temp_dir:
mp = _model.ModelPipeline('wf', 'dv', civisml_version=version)
# Before pickling, make sure the template IDs are set as expected
assert mp.train_template_id == train_id
assert mp.predict_template_id == predict_id
pickle_path = os.path.join(temp_dir, 'model.pkl')
with open(pickle_path, 'wb') as f:
pickle.dump(mp, f)
with open(pickle_path, 'rb') as f:
mp_unpickled = pickle.load(f)
# After unpickling, the template IDs should remain.
assert mp_unpickled.train_template_id == train_id
assert mp_unpickled.predict_template_id == predict_id
def mp_setup():
mock_api = setup_client_mock()
mp = _model.ModelPipeline('wf', 'dv', client=mock_api)
return mp
def mp_setup():
mock_api = create_client_mock_for_container_tests()
mock_api.aliases.list.return_value = TEST_TEMPLATE_ID_ALIAS_OBJECTS
mp = _model.ModelPipeline('wf', 'dv', client=mock_api)
return mp
def mp_setup():
mock_api = setup_client_mock()
mock_api.aliases.list.return_value = TEST_TEMPLATE_ID_ALIAS_OBJECTS
mp = _model.ModelPipeline('wf', 'dv', client=mock_api)
return mp
