def test_send_to_bq(self):
payloads = []
def json_request(url, method, payload, scopes, deadline):
self.assertEqual(
'https://www.googleapis.com/bigquery/v2/projects/sample-app/datasets/'
'swarming/tables/foo/insertAll',
url)
payloads.append(payload)
self.assertEqual('POST', method)
self.assertEqual(bq_state.bqh.INSERT_ROWS_SCOPE, scopes)
self.assertEqual(600, deadline)
return {'insertErrors': []}
self.mock(bq_state.net, 'json_request', json_request)
rows = [
('key1', struct_pb2.Struct()),
('key2', struct_pb2.Struct()),
]
self.assertEqual(0, bq_state.send_to_bq('foo', rows))
expected = [
{
'ignoreUnknownValues': False,
'kind': 'bigquery#tableDataInsertAllRequest',
'skipInvalidRows': True,
},
]
actual_rows = payloads[0].pop('rows')
self.assertEqual(expected, payloads)
self.assertEqual(2, len(actual_rows))
def test_send_to_bq_fail(self):
# Test the failure code path.
payloads = []
def json_request(url, method, payload, scopes, deadline):
self.assertEqual(
'https://www.googleapis.com/bigquery/v2/projects/sample-app/datasets/'
'swarming/tables/foo/insertAll', url)
first = not payloads
payloads.append(payload)
self.assertEqual('POST', method)
self.assertEqual(bq_state.bqh.INSERT_ROWS_SCOPE, scopes)
self.assertEqual(600, deadline)
# Return an error on the first call.
if first:
return {
'insertErrors': [
{
'index':
0,
'errors': [
{
'reason': 'sadness',
'message': 'Oh gosh',
},
],
},
],
}
return {'insertErrors': []}
self.mock(bq_state.net, 'json_request', json_request)
rows = [
('key1', struct_pb2.Struct()),
('key2', struct_pb2.Struct()),
]
self.assertEqual(1, bq_state.send_to_bq('foo', rows))
self.assertEqual(2, len(payloads), payloads)
expected = {
'ignoreUnknownValues': False,
'kind': 'bigquery#tableDataInsertAllRequest',
'skipInvalidRows': True,
}
actual_rows = payloads[0].pop('rows')
self.assertEqual(expected, payloads[0])
self.assertEqual(2, len(actual_rows))
expected = {
'ignoreUnknownValues': False,
'kind': 'bigquery#tableDataInsertAllRequest',
'skipInvalidRows': True,
}
actual_rows = payloads[1].pop('rows')
self.assertEqual(expected, payloads[1])
self.assertEqual(1, len(actual_rows))
def task_bq_events(start, end):
"""Sends BotEvents to BigQuery swarming.bot_events table."""
def _convert(e):
"""Returns a tuple(bq_key, row)."""
out = swarming_pb2.BotEvent()
e.to_proto(out)
bq_key = e.id + ':' + e.ts.strftime(u'%Y-%m-%dT%H:%M:%S.%fZ')
return (bq_key, out)
total = 0
failed = 0
q = BotEvent.query(BotEvent.ts >= start, BotEvent.ts <= end)
cursor = None
more = True
while more:
entities, cursor, more = q.fetch_page(500, start_cursor=cursor)
total += len(entities)
failed += bq_state.send_to_bq('bot_events',
[_convert(e) for e in entities])
return total, failed
def test_send_to_bq_empty(self):
# Empty, nothing is done. No need to mock the HTTP client.
self.assertEqual(0, bq_state.send_to_bq('foo', []))