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))