def check_estimated_size_bytes(self, entity_bytes, timestamp, namespace=None):
"""A helper method to test get_estimated_size_bytes"""
self._mock_client.namespace = namespace
self._mock_client.query.return_value = self._mock_query
self._mock_query.project = self._PROJECT
self._mock_query.namespace = namespace
self._mock_query.fetch.side_effect = [
[{'timestamp': timestamp}],
[{'entity_bytes': entity_bytes}],
]
self._mock_query.kind = self._KIND
split_query_fn = ReadFromDatastore._SplitQueryFn(num_splits=0)
self.assertEqual(entity_bytes,
split_query_fn.get_estimated_size_bytes(self._mock_client,
self._mock_query))
if namespace is None:
ns_keyword = '_'
else:
ns_keyword = '_Ns_'
self._mock_client.query.assert_has_calls([
call(kind='__Stat%sTotal__' % ns_keyword, order=['-timestamp']),
call().fetch(limit=1),
call(kind='__Stat%sKind__' % ns_keyword),
call().add_filter('kind_name', '=', self._KIND),
call().add_filter('timestamp', '=', timestamp),
call().fetch(limit=1),
])
def test_SplitQueryFn_with_query_limit(self):
"""A test that verifies no split is performed when the query has a limit."""
with patch.object(helper, 'get_client', return_value=self._mock_client):
num_splits = 4
expected_num_splits = 1
self._mock_query.limit = 3
split_query_fn = ReadFromDatastore._SplitQueryFn(num_splits)
split_queries = split_query_fn.process(self._mock_query)
self.assertEqual(expected_num_splits, len(split_queries))
def test_SplitQueryFn_with_num_splits(self):
with patch.object(helper, 'get_client', return_value=self._mock_client):
num_splits = 23
expected_num_splits = 23
def fake_get_splits(unused_client, query, num_splits):
return [query] * num_splits
with patch.object(query_splitter, 'get_splits',
side_effect=fake_get_splits):
split_query_fn = ReadFromDatastore._SplitQueryFn(num_splits)
split_queries = split_query_fn.process(self._mock_query)
self.assertEqual(expected_num_splits, len(split_queries))
def test_SplitQueryFn_with_exception(self):
"""A test that verifies that no split is performed when failures occur."""
with patch.object(helper, 'get_client', return_value=self._mock_client):
# Force _SplitQueryFn to compute the number of query splits
num_splits = 0
expected_num_splits = 1
entity_bytes = (expected_num_splits *
ReadFromDatastore._DEFAULT_BUNDLE_SIZE_BYTES)
with patch.object(
ReadFromDatastore._SplitQueryFn, 'get_estimated_size_bytes',
return_value=entity_bytes):
with patch.object(query_splitter, 'get_splits',
side_effect=query_splitter.QuerySplitterError(
"Testing query split error")):
split_query_fn = ReadFromDatastore._SplitQueryFn(num_splits)
split_queries = split_query_fn.process(self._mock_query)
self.assertEqual(expected_num_splits, len(split_queries))
self.assertEqual(self._mock_query, split_queries[0])
def test_SplitQueryFn_without_num_splits(self):
with patch.object(helper, 'get_client', return_value=self._mock_client):
# Force _SplitQueryFn to compute the number of query splits
num_splits = 0
expected_num_splits = 23
entity_bytes = (expected_num_splits *
ReadFromDatastore._DEFAULT_BUNDLE_SIZE_BYTES)
with patch.object(
ReadFromDatastore._SplitQueryFn, 'get_estimated_size_bytes',
return_value=entity_bytes):
def fake_get_splits(unused_client, query, num_splits):
return [query] * num_splits
with patch.object(query_splitter, 'get_splits',
side_effect=fake_get_splits):
split_query_fn = ReadFromDatastore._SplitQueryFn(num_splits)
split_queries = split_query_fn.process(self._mock_query)
self.assertEqual(expected_num_splits, len(split_queries))