def _build_mock_client_for_copying(self, table_exists=True):
bq_client = mock.MagicMock()
if not table_exists:
bq_client.get_table.side_effect = exceptions.NotFound('not found')
bq = GoogleBigQuery()
bq._client = bq_client
return bq
def _build_mock_client_for_querying(self, results):
# Create a mock that will play the role of the query job
query_job = mock.MagicMock()
query_job.result.return_value = FakeResults(results)
# Create a mock that will play the role of our GoogleBigQuery client
client = mock.MagicMock()
client.query.return_value = query_job
bq = GoogleBigQuery()
bq._client = client
return bq
def test_query__no_results(self):
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = 'foo'
query_string = 'select * from table'
# Pass the mock class into our GoogleBigQuery constructor
bq = GoogleBigQuery()
bq._client = self._build_mock_client([])
# Run a query against our parsons GoogleBigQuery class
result = bq.query(query_string)
# Check our return value
self.assertEqual(result, None)
def test_query(self):
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = 'foo'
query_string = 'select * from table'
# Pass the mock class into our GoogleBigQuery constructor
bq = GoogleBigQuery()
bq._client = self._build_mock_client([{'one': 1, 'two': 2}])
# Run a query against our parsons GoogleBigQuery class
result = bq.query(query_string)
# Check our return value
self.assertEqual(result.num_rows, 1)
self.assertEqual(result.columns, ['one', 'two'])
self.assertEqual(result[0], {'one': 1, 'two': 2})
def _build_mock_client_for_querying(self, results):
# Create a mock that will play the role of the cursor
cursor = mock.MagicMock()
cursor.execute.return_value = None
cursor.fetchmany.side_effect = [results, []]
# Create a mock that will play the role of the connection
connection = mock.MagicMock()
connection.cursor.return_value = cursor
# Create a mock that will play the role of the Google BigQuery dbapi module
dbapi = mock.MagicMock()
dbapi.connect.return_value = connection
# Create a mock that will play the role of our GoogleBigQuery client
client = mock.MagicMock()
bq = GoogleBigQuery()
bq._client = client
bq._dbapi = dbapi
return bq