def test_modelpipeline_train(mock_ccr, mp_setup):
mp = mp_setup
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'
def test_modelpipeline_train_file_name(mock_ccr, mock_stash, mp_setup):
mp = mp_setup
mock1, mock2 = mock.Mock(), mock.Mock()
mock_ccr.return_value = 'res', mock1, mock2
assert 'res' == mp.train(csv_path='meaning_of_life.csv')
mock_stash.assert_called_once_with('meaning_of_life.csv', client=mock.ANY)
assert mp.train_result_ == 'res'
def setup_listener_status_mocks(status_category):
match = mock.Mock()
callback = mock.Mock()
disconnect = mock.Mock()
listener = JobCompleteListener(match, callback, disconnect)
status = mock.Mock()
status.category = status_category
return match, callback, disconnect, listener, status
def test_modelpipeline_train_df(mock_ccr, mock_stash, mp_setup):
mp = mp_setup
mock1, mock2 = mock.Mock(), mock.Mock()
mock_ccr.return_value = 'res', mock1, mock2
train_data = pd.DataFrame({'a': [1, 2], 'b': [3, 4]})
assert 'res' == mp.train(train_data)
mock_stash.assert_called_once_with(train_data, client=mock.ANY)
assert mp.train_result_ == 'res'
def test_store_last_response():
mock_client = mock.Mock()
endpoint = Endpoint({}, client=mock_client, return_type='raw')
returned_resp = {'value': 'response'}
endpoint._make_request = mock.Mock(return_value=returned_resp)
resp = endpoint._call_api('GET')
assert resp == returned_resp
assert mock_client.last_response is resp
def test_listener_calls_callback_when_message_matches(self):
match = mock.Mock()
match.return_value = True
callback = mock.Mock()
listener = JobCompleteListener(match, callback)
message = mock.Mock()
message.message.return_value = 'test message'
listener.message(None, message)
match.assert_called_with(message.message)
self.assertEqual(callback.call_count, 1)
def test_listener_does_not_call_callback(self):
match = mock.Mock()
match.return_value = False
callback = mock.Mock()
listener = JobCompleteListener(match, callback)
message = mock.Mock()
message.message.return_value = 'test message'
listener.message(None, message)
match.assert_called_with(message.message)
self.assertEqual(callback.call_count, 0)
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_set_api_result_result_succeeded(self, mock_subscribe, mock_api):
mock_pubnub = mock.Mock()
mock_pubnub.unsubscribe_all.return_value = None
mock_subscribe.return_value = mock_pubnub
poller = mock.Mock()
api_result = mock.Mock()
api_result.state = 'succeeded'
result = CivisFuture(poller, (1, 2))
result._set_api_result(api_result)
assert poller.call_count == 0
assert mock_pubnub.unsubscribe_all.call_count == 1
assert result._state == 'FINISHED'
def test_set_api_result_failed(self, mock_subscribe, mock_api):
mock_pubnub = mock.Mock()
mock_pubnub.unsubscribe_all.return_value = None
mock_subscribe.return_value = mock_pubnub
poller = mock.Mock()
api_result = mock.Mock()
api_result.state = 'failed'
result = CivisFuture(poller, (1, 2))
result._set_api_result(api_result)
assert mock_pubnub.unsubscribe_all.call_count == 1
assert result._state == 'FINISHED'
with pytest.raises(CivisJobFailure):
result.result()
def test_file_to_civis(mock_open, mock_file_to_civis_helper):
# Test that passing a path to file_to_civis opens a file.
civis.io.file_to_civis("foo", "foo_name")
mock_open.return_value = mock_file = mock.Mock()
assert mock_open.called_once_with("foo", "rb")
assert mock_file_to_civis_helper.called_once_with("foo", "foo_name",
mock_file)
def test_modelpipeline_classmethod_constructor(mock_future,
container_response_stub):
mock_client = mock.Mock()
mock_client.scripts.get_containers.return_value = \
container = container_response_stub
mock_client.credentials.get.return_value = Response({'name': 'Token'})
resources = {
'REQUIRED_CPU': 1000,
'REQUIRED_MEMORY': 9999,
'REQUIRED_DISK_SPACE': -20
}
# test everything is working fine
mp = _model.ModelPipeline.from_existing(1, 1, client=mock_client)
assert isinstance(mp, _model.ModelPipeline)
assert mp.dependent_variable == [container.arguments['TARGET_COLUMN']]
assert mp.primary_key == container.arguments['PRIMARY_KEY']
excluded = container.arguments.get('EXCLUDE_COLS', None)
assert mp.excluded_columns == excluded.split() if excluded else None
assert mp.model == container.arguments['MODEL']
assert mp.calibration == container.arguments['CALIBRATION']
assert mp.cv_params == json.loads(container.arguments['CVPARAMS'])
assert mp.parameters == json.loads(container.arguments['PARAMS'])
assert mp.job_resources == resources
assert mp.model_name == container.name[:-6]
assert mp.notifications == {
camel_to_snake(key): val
for key, val in container.notifications.items()
}
deps = container.arguments.get('DEPENDENCIES', None)
assert mp.dependencies == deps.split() if deps else None
assert mp.git_token_name == 'Token'
def test_file_to_dataframe_infer_gzip():
m_client = mock.Mock()
m_client.files.get.return_value = Response({'name': 'spam.csv.gz',
'file_url': 'url'})
with mock.patch.object(civis.io._files.pd, 'read_csv') as mock_read_csv:
civis.io.file_to_dataframe(121, compression='infer', client=m_client)
assert mock_read_csv.called_once_with(121, compression='gzip')
def test_file_id_from_run_output_platform_error():
# Make sure we don't swallow real Platform errors
m_client = mock.Mock()
m_client.scripts.list_containers_runs_outputs.side_effect =\
MockAPIError(500) # Mock a platform error
with pytest.raises(CivisAPIError):
civis.io.file_id_from_run_output('name', 17, 13, client=m_client)
def test_stash_local_data_from_dataframe_csv(mock_file):
df = pd.DataFrame({'a': [1], 'b': [2]})
assert _model._stash_dataframe_as_csv(df, mock.Mock()) == -11
mock_file.assert_called_once_with(mock.ANY,
name='modelpipeline_data.csv',
client=mock.ANY)
assert isinstance(mock_file.call_args[0][0], StringIO)
def test_modelpipeline_classmethod_constructor_errors(mock_future):
# catch 404 error if model isn't found and throw ValueError
mock_client = mock.Mock()
response = namedtuple('Reponse', ['content', 'response', 'reason',
'status_code'])(False, None, None, 404)
mock_future.side_effect = CivisAPIError(response)
with pytest.raises(ValueError):
_model.ModelPipeline.from_existing(1, 1, client=mock_client)
def test_load_json(mock_c2f):
obj = {'spam': 'eggs'}
def _dump_json(file_id, buf, *args, **kwargs):
buf.write(json.dumps(obj).encode())
mock_c2f.side_effect = _dump_json
out = civis.io.file_to_json(13, client=mock.Mock())
assert out == obj
def test_get_table_id(schema_tablename):
"""Check that get_table_id handles quoted schema.tablename correctly."""
client = civis.APIClient(local_api_spec=TEST_SPEC, api_key='none')
client.get_database_id = mock.Mock(return_value=123)
mock_tables = mock.MagicMock()
mock_tables.__getitem__.side_effect = {0: mock.Mock()}.__getitem__
client.tables.list = mock.Mock(return_value=mock_tables)
client.get_table_id(table=schema_tablename, database=123)
client.tables.list.assert_called_once_with(
database_id=123,
schema='foo',
name='bar'
)
def test_file_id_from_run_output_exact():
m_client = mock.Mock()
m_client.scripts.list_containers_runs_outputs.return_value = \
[Response({'name': 'spam', 'object_id': 2013,
'object_type': 'File'})]
fid = civis.io.file_id_from_run_output('spam', 17, 13, client=m_client)
assert fid == 2013
def test_file_to_dataframe_kwargs():
m_client = mock.Mock()
m_client.files.get.return_value = Response({'name': 'spam.csv',
'file_url': 'url'})
with mock.patch.object(civis.io._files.pd, 'read_csv') as mock_read_csv:
civis.io.file_to_dataframe(121, compression='special', client=m_client,
delimiter='|', nrows=10)
assert mock_read_csv.called_once_with(121, compression='special',
delimiter='|', nrows=10)
def test_file_id_from_run_output_no_run():
# Get an IOError if we request a file from a run which doesn't exist
m_client = mock.Mock()
m_client.scripts.list_containers_runs_outputs.side_effect =\
MockAPIError(404) # Mock a run which doesn't exist
with pytest.raises(IOError) as err:
civis.io.file_id_from_run_output('name', 17, 13, client=m_client)
assert 'could not find job/run id 17/13' in str(err.value).lower()
def test_poll_on_creation(self):
poller = mock.Mock(return_value=Response({"state": "running"}))
pollable = PollableResult(poller, (),
polling_interval=0.01,
poll_on_creation=False)
pollable.done() # Check status once to start the polling thread
assert poller.call_count == 0
time.sleep(0.015)
assert poller.call_count == 1
def test_poller_returns_none(self):
poller = mock.Mock(
side_effect=[None, None,
Response({'state': 'success'})])
polling_thread = _ResultPollingThread(poller, (),
polling_interval=0.01)
polling_thread.run()
time.sleep(0.05)
assert poller.call_count == 3
def test_modelpipeline_estimator(mp_setup):
mp = mp_setup
with pytest.raises(ValueError,
message="This model hasn't been trained yet."):
mp.estimator
mp.train_result_ = mock.Mock()
mp.train_result_.estimator = 'foo'
assert mp.estimator == 'foo'
def test_file_to_dataframe_infer():
m_client = mock.Mock()
url = 'url'
m_client.files.get.return_value = Response({'name': 'spam.csv',
'file_url': url})
with mock.patch.object(civis.io._files.pd, 'read_csv',
autospec=True) as mock_read_csv:
civis.io.file_to_dataframe(121, compression='infer', client=m_client)
mock_read_csv.assert_called_once_with(url, compression='infer')
def test_file_id_from_run_output_approximate():
# Test fuzzy name matching
m_client = mock.Mock()
m_client.scripts.list_containers_runs_outputs.return_value = \
[Response({'name': 'spam.csv.gz', 'object_id': 2013,
'object_type': 'File'})]
fid = civis.io.file_id_from_run_output('^spam', 17, 13, regex=True,
client=m_client)
assert fid == 2013
def test_file_id_from_run_output_no_file():
# Get an IOError if we request a file which doesn't exist
m_client = mock.Mock()
m_client.scripts.list_containers_runs_outputs.return_value = [
Response({'name': 'spam', 'object_id': 2013,
'object_type': 'File'})]
with pytest.raises(FileNotFoundError) as err:
civis.io.file_id_from_run_output('eggs', 17, 13, client=m_client)
assert 'not an output' in str(err.value)
def setup_client_mock(script_id=-10,
run_id=100,
state='succeeded',
run_outputs=None):
"""Return a Mock set up for use in testing container scripts
Parameters
----------
script_id: int
Mock-create containers with this ID when calling `post_containers`
or `post_containers_runs`.
run_id: int
Mock-create runs with this ID when calling `post_containers_runs`.
state: str, optional
The reported state of the container run
run_outputs: list, optional
List of Response objects returned as run outputs
Returns
-------
`unittest.mock.Mock`
With scripts endpoints `post_containers`, `post_containers_runs`,
`post_cancel`, and `get_containers_runs` set up.
"""
c = mock.Mock()
c.__class__ = APIClient
mock_container = Response({'id': script_id})
c.scripts.post_containers.return_value = mock_container
mock_container_run_start = Response({
'id': run_id,
'container_id': script_id,
'state': 'queued'
})
mock_container_run = Response({
'id': run_id,
'container_id': script_id,
'state': state
})
c.scripts.post_containers_runs.return_value = mock_container_run_start
c.scripts.get_containers_runs.return_value = mock_container_run
c.scripts.list_containers_runs_outputs.return_value = (run_outputs or [])
c.scripts.list_containers_runs_logs.return_value = []
def change_state_to_cancelled(script_id):
mock_container_run.state = "cancelled"
return mock_container_run
c.scripts.post_cancel.side_effect = change_state_to_cancelled
# Avoid channels endpoint while testing here
del c.channels
return c
def test_modelpipeline_classmethod_constructor_defaults(
mock_future, container_response_stub):
del container_response_stub.arguments['PARAMS']
del container_response_stub.arguments['CVPARAMS']
mock_client = mock.Mock()
mock_client.scripts.get_containers.return_value = container_response_stub
mock_client.credentials.get.return_value = Response({'name': 'Token'})
# test everything is working fine
mp = _model.ModelPipeline.from_existing(1, 1, client=mock_client)
assert mp.cv_params == {}
assert mp.parameters == {}
