def test_sql_script():
sql = "SELECT SPECIAL SQL QUERY"
export_job_id = 32
database_id = 29
credential_id = 3920
response = Response({'id': export_job_id})
mock_client = create_client_mock()
mock_client.scripts.post_sql.return_value = response
mock_client.get_database_id.return_value = database_id
mock_client.default_credential = credential_id
civis.io._tables._sql_script(client=mock_client,
sql=sql,
database='fake-db',
job_name='My job',
credential_id=None,
hidden=False,
csv_settings=None)
mock_client.scripts.post_sql.assert_called_once_with(
'My job',
remote_host_id=database_id,
credential_id=credential_id,
sql=sql,
hidden=False,
csv_settings={})
mock_client.scripts.post_sql_runs.assert_called_once_with(export_job_id)
def test_list_models():
resp = [Response({'id': 2834, 'name': 'RFC model'})]
m_client = create_client_mock()
m_client.scripts.list_custom.return_value = resp
out = list_models(job_type='train', client=m_client)
assert out == resp
out = list_models(job_type='predict', client=m_client)
assert out == resp
out = list_models(job_type=None, client=m_client)
assert out == resp
def test_civis_to_file_local(mock_requests):
# Test that a call to civis_to_file uses `requests` to grab the contents
# of a URL given by the API client and writes it to a file.
mock_civis = create_client_mock()
mock_requests.get.return_value.iter_content.return_value =\
(l.encode() for l in 'abcdef')
with TemporaryDirectory() as tdir:
fname = os.path.join(tdir, 'testfile')
_files.civis_to_file(137, fname, client=mock_civis)
with open(fname, 'rt') as _fin:
assert _fin.read() == 'abcdef'
mock_civis.files.get.assert_called_once_with(137)
mock_requests.get.assert_called_once_with(
mock_civis.files.get.return_value.file_url, stream=True)
def _create_share_model_client_mock(run_ids):
m_client = create_client_mock()
m_client.scripts.put_containers_shares_users.return_value = 'usershare'
m_client.scripts.put_containers_shares_groups.return_value = 'groupshare'
m_client.scripts.delete_containers_shares_users.return_value = 'userdel'
m_client.scripts.delete_containers_shares_groups.return_value = 'groupdel'
m_client.scripts.list_containers_runs.return_value = [
Response({'id': _id}) for _id in run_ids]
m_client.scripts.list_containers_runs_outputs.return_value = [
Response({'object_id': 117, 'object_type': 'File', 'name': 'fname'}),
Response({'object_id': 31, 'object_type': 'Project'}),
Response({'object_id': 37, 'object_type': 'JSONValue'}),
]
return m_client
def test_file_to_civis(mock_requests, input_filename):
# Test that file_to_civis posts a Civis File with the API client
# and calls `requests.post` on the returned URL.
mock_civis = create_client_mock()
civis_name, expected_id = 'newname', 137
mock_civis.files.post.return_value.id = expected_id
with TemporaryDirectory() as tdir:
fname = os.path.join(tdir, 'newname')
with open(fname, 'wt') as _fout:
_fout.write('abcdef')
fid = _files.file_to_civis(fname, input_filename, expires_at=None,
client=mock_civis)
assert fid == expected_id
mock_civis.files.post.assert_called_once_with(civis_name, expires_at=None)
mock_requests.post.assert_called_once_with(
mock_civis.files.post.return_value.upload_url, files=mock.ANY)
def create_mock_client_with_job():
mock_client = create_client_mock()
job_response = Response({"id": 1, "name": "test"})
run_post_response = Response({"id": 1})
run_get_response = Response({
"id": run_post_response.id,
"state": "succeeded",
"is_cancel_requested": False,
"error": None,
"custom_id": run_post_response.id,
})
mock_client.scripts.post_custom.return_value = job_response
mock_client.scripts.post_custom_runs.return_value = run_post_response
mock_client.scripts.get_custom_runs.return_value = run_get_response
return mock_client
def test_civis_to_file_retries(mock_requests):
mock_civis = create_client_mock()
# Mock the request iter_content so it fails partway the first time.
# Python 2.7 doesn't have the nonlocal keyword, so here's a little class to
# track whether it's the first call.
class UnreliableIterContent:
def __init__(self):
self.first_try = True
def mock_iter_content(self, _):
chunks = [l.encode() for l in 'abcdef']
for i, chunk in enumerate(chunks):
# Fail partway through on the first try.
if self.first_try and i == 3:
self.first_try = False
raise requests.ConnectionError()
yield chunk
mock_requests.get.return_value.iter_content = \
UnreliableIterContent().mock_iter_content
# Add some data to the buffer to test that we seek to the right place
# when retrying.
buf = io.BytesIO(b'0123')
buf.seek(4)
_files.civis_to_file(137, buf, client=mock_civis)
# Check that retries work and that the buffer position is reset.
# If we didn't seek when retrying, we'd get abcabcdef.
# If we seek'd to position 0, then we'd get abcdef.
buf.seek(0)
assert buf.read() == b'0123abcdef'
mock_civis.files.get.assert_called_once_with(137)
assert mock_requests.get.call_count == 2
mock_requests.get.assert_called_with(
mock_civis.files.get.return_value.file_url, stream=True)
def test_export_to_civis_file(mock_sql_script):
expected = [{'file_id': 9844453}]
mock_client = create_client_mock()
response = Response({'state': 'success', 'output': expected})
mock_client.scripts.get_sql_runs.return_value = response
mock_client.scripts.post_sql
sql = "SELECT 1"
fut = civis.io.export_to_civis_file(sql, 'fake-db',
polling_interval=POLL_INTERVAL,
client=mock_client)
data = fut.result()['output']
assert data == expected
mock_sql_script.assert_called_once_with(client=mock_client,
sql=sql,
database='fake-db',
job_name=None,
credential_id=None,
csv_settings=None,
hidden=True)
def test_civis_to_file_retries(mock_requests):
mock_civis = create_client_mock()
first_try = True
# Mock the request iter_content so it fails partway the first time.
def mock_iter_content(_):
nonlocal first_try
chunks = [l.encode() for l in 'abcdef']
for i, chunk in enumerate(chunks):
# Fail partway through on the first try.
if first_try and i == 3:
first_try = False
raise requests.ConnectionError()
yield chunk
mock_requests.get.return_value.iter_content = mock_iter_content
# Add some data to the buffer to test that we seek to the right place
# when retrying.
buf = io.BytesIO(b'0123')
buf.seek(4)
_files.civis_to_file(137, buf, client=mock_civis)
# Check that retries work and that the buffer position is reset.
# If we didn't seek when retrying, we'd get abcabcdef.
# If we seek'd to position 0, then we'd get abcdef.
buf.seek(0)
assert buf.read() == b'0123abcdef'
mock_civis.files.get.assert_called_once_with(137)
assert mock_requests.get.call_count == 2
mock_requests.get.assert_called_with(
mock_civis.files.get.return_value.file_url, stream=True)
def client_mock(self):
self.mock_client = create_client_mock()
def test_client_mock_attributeerror():
mock_client = mocks.create_client_mock()
with pytest.raises(AttributeError):
mock_client.not_an_endpoint()
def test_client_mock_bad_parameter():
mock_client = mocks.create_client_mock()
mock_client.tables.list(database_id=1) # Valid parameter
with pytest.raises(TypeError):
mock_client.tables.list(db_id=1) # Invalid parameter
