def testBigqueryPagedResults(self, _):
mock_client = mock.Mock()
mock_client.jobs().query().execute.return_value = {
"totalRows": 2,
"jobComplete": True,
"schema": {
"fields": 'mock_fields'
},
"rows": ["rows"],
"pageToken": "dummy page token",
"jobReference": {
"jobId": "dummy job id"
},
}
mock_client.jobs().getQueryResults().execute.return_value = {
"schema": {
"fields": 'mock_fields'
},
"rows": ["rows2"],
}
success, rows = bigquery_helper.QueryRequest(mock_client, 'project',
'query')
self.assertTrue(mock_client.jobs().query().execute.called)
self.assertTrue(mock_client.jobs().getQueryResults().execute.called)
self.assertTrue(success)
self.assertEqual(len(rows), 2)
def testBigqueryQueryRequestWithMissingFields(self):
mock_client = mock.Mock()
mock_client.jobs().query().execute.return_value = {}
success, rows = bigquery_helper.QueryRequest(mock_client, 'project',
'query')
self.assertTrue(mock_client.jobs().query().execute.called)
self.assertFalse(success)
self.assertEqual(rows, [])
def testBigqueryQueryRequestNoRows(self):
mock_client = mock.Mock()
mock_client.jobs().query().execute.return_value = {
'jobComplete': True,
'totalRows': '0'
}
success, rows = bigquery_helper.QueryRequest(mock_client, 'project',
'query')
self.assertTrue(mock_client.jobs().query().execute.called)
self.assertTrue(success)
self.assertEqual(rows, [])
def testBigqueryQueryRequest(self, _):
mock_client = mock.Mock()
mock_client.jobs().query().execute.return_value = {
"totalRows": 1,
"jobComplete": True,
"schema": {
"fields": 'mock_fields'
},
"rows": ["rows"]
}
success, rows = bigquery_helper.QueryRequest(mock_client, 'project',
'query')
self.assertTrue(mock_client.jobs().query().execute.called)
self.assertTrue(success)
self.assertEqual(len(rows), 1)