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