def testSplitQueryTestFilterInitialSplit(self): """Tests that initial query splitting works as expected.""" test_filter = queries._SplitQueryTestFilter(['1', '2', '3'], 2) self.assertEqual(test_filter._test_id_lists, [['1', '2'], ['3']]) self.assertEqual(len(test_filter.GetClauses()), 2) test_filter = queries._SplitQueryTestFilter(['1', '2', '3'], 3) self.assertEqual(test_filter._test_id_lists, [['1', '2', '3']]) self.assertEqual(len(test_filter.GetClauses()), 1)
def testRetryOnMemoryLimit(self):
"""Tests that queries are split and retried if the memory limit is hit."""
def SideEffect(*_, **__):
SideEffect.call_count += 1
if SideEffect.call_count == 1:
raise queries.MemoryLimitError()
return []
SideEffect.call_count = 0
with mock.patch.object(
self._querier,
'_GetTestFilterForBuilder',
return_value=queries._SplitQueryTestFilter(
['filter_a', 'filter_b'], 10)), mock.patch.object(
self._querier,
'_RunBigQueryCommandsForJsonOutput') as query_mock:
query_mock.side_effect = SideEffect
self._querier.QueryBuilder('builder', 'ci')
self.assertEqual(query_mock.call_count, 2)
args, _ = unittest_utils.GetArgsForMockCall(query_mock.call_args_list, 0)
first_query = args[0][0]
self.assertIn('filter_a', first_query)
self.assertIn('filter_b', first_query)
args, _ = unittest_utils.GetArgsForMockCall(query_mock.call_args_list, 1)
second_query_first_half = args[0][0]
self.assertIn('filter_a', second_query_first_half)
self.assertNotIn('filter_b', second_query_first_half)
second_query_second_half = args[0][1]
self.assertIn('filter_b', second_query_second_half)
self.assertNotIn('filter_a', second_query_second_half)
def testSplitQueryTestFilterSplitQueryCannotSplitFurther(self):
"""Tests that SplitQueryTestFilter's failure mode."""
test_filter = queries._SplitQueryTestFilter(['1'], 1)
with self.assertRaises(queries.QuerySplitError):
test_filter.SplitFilter()
def testSplitQueryTestFilterSplitQuery(self): """Tests that SplitQueryTestFilter's query splitting works.""" test_filter = queries._SplitQueryTestFilter(['1', '2'], 10) self.assertEqual(len(test_filter.GetClauses()), 1) test_filter.SplitFilter() self.assertEqual(len(test_filter.GetClauses()), 2)