def test_refresh_token_raises_cartoexception(self):
refresh_token_checker = RefreshTokenChecker('', 0)
original_query_method = bigquery.Client.query
bigquery.Client.query = refresh_token_checker.query_raiser
bq_client = BigQueryClient(self.credentials)
with pytest.raises(DOError):
bq_client.query('select * from')
bigquery.Client.query = original_query_method
def test_refresh_token(self):
expected_response = 'ok'
refresh_token_checker = RefreshTokenChecker(expected_response, 2)
original_query_method = bigquery.Client.query
bigquery.Client.query = refresh_token_checker.query_raiser
bq_client = BigQueryClient(self.credentials)
response = bq_client.query('select * from')
assert response == expected_response
bigquery.Client.query = original_query_method
def test_download_using_if_exists(self):
project = _WORKING_PROJECT
dataset = 'fake_dataset'
table = 'fake_table'
file_path = self.file_path
bq_client = BigQueryClient(self.credentials)
query = 'SELECT * FROM `{}.{}.{}`'.format(project, dataset, table)
job = bq_client.query(query)
with open(file_path, 'w'):
with self.assertRaises(OSError):
bq_client.download_to_file(job, file_path, fail_if_exists=True, progress_bar=False)
def test_download_to_dataframe_full(self, download_mock,
column_names_mock):
data = [{'column1': 'word', 'column2': 'word word'}]
columns = ['column1', 'column2']
column_names_mock.return_value = Mock(return_value=columns)
download_mock.return_value = data
expected_df = pd.DataFrame(data, columns=columns)
bq_client = BigQueryClient(self.credentials)
job = QueryJobMock(data)
df = bq_client.download_to_dataframe(job)
assert df.equals(expected_df)
def test_download_to_file_full(self, download_mock, column_names_mock):
data = [{'0': 'word', '1': 'word word'}]
columns = ['column1', 'column2']
column_names_mock.return_value = Mock(return_value=columns)
download_mock.return_value = data
file_path = self.file_path
bq_client = BigQueryClient(self.credentials)
job = QueryJobMock(data)
bq_client.download_to_file(job,
file_path,
column_names=columns,
progress_bar=False)
rows = []
with open(file_path) as csvfile:
csvreader = csv.reader(csvfile)
rows.append(next(csvreader))
rows.append(next(csvreader))
assert rows[0] == columns
assert rows[1] == list(data[0].values())
def test_instantiation(self):
bq_client = BigQueryClient(self.credentials)
assert isinstance(bq_client, BigQueryClient)