diff --git a/datastore/google/cloud/datastore/_http.py b/datastore/google/cloud/datastore/_http.py index a4c0a3f8c906..ac9059ff0340 100644 --- a/datastore/google/cloud/datastore/_http.py +++ b/datastore/google/cloud/datastore/_http.py @@ -379,10 +379,8 @@ def run_query(self, project, query_pb, namespace=None, the given transaction. Incompatible with ``eventual==True``. - :rtype: tuple - :returns: Four-tuple containing the entities returned, - the end cursor of the query, a ``more_results`` - enum and a count of the number of skipped results. + :rtype: :class:`.datastore_pb2.RunQueryResponse` + :returns: The protobuf response from a ``runQuery`` request. """ request = _datastore_pb2.RunQueryRequest() _set_read_options(request, eventual, transaction_id) @@ -391,13 +389,7 @@ def run_query(self, project, query_pb, namespace=None, request.partition_id.namespace_id = namespace request.query.CopyFrom(query_pb) - response = self._datastore_api.run_query(project, request) - return ( - [e.entity for e in response.batch.entity_results], - response.batch.end_cursor, # Assume response always has cursor. - response.batch.more_results, - response.batch.skipped_results, - ) + return self._datastore_api.run_query(project, request) def begin_transaction(self, project): """Begin a transaction. diff --git a/datastore/google/cloud/datastore/query.py b/datastore/google/cloud/datastore/query.py index e8989a41a9dd..9ee565f2f2f2 100644 --- a/datastore/google/cloud/datastore/query.py +++ b/datastore/google/cloud/datastore/query.py @@ -441,43 +441,33 @@ def _build_protobuf(self): return pb - def _process_query_results(self, entity_pbs, cursor_as_bytes, - more_results_enum, skipped_results): + def _process_query_results(self, response_pb): """Process the response from a datastore query. - :type entity_pbs: iterable - :param entity_pbs: The entities returned in the current page. - - :type cursor_as_bytes: bytes - :param cursor_as_bytes: The end cursor of the query. - - :type more_results_enum: - :class:`.query_pb2.QueryResultBatch.MoreResultsType` - :param more_results_enum: Enum indicating if there are more results. - - :type skipped_results: int - :param skipped_results: The number of skipped results. + :type response_pb: :class:`.datastore_pb2.RunQueryResponse` + :param response_pb: The protobuf response from a ``runQuery`` request. :rtype: iterable :returns: The next page of entity results. :raises ValueError: If ``more_results`` is an unexpected value. """ - self._skipped_results = skipped_results + self._skipped_results = response_pb.batch.skipped_results - if cursor_as_bytes == b'': # Empty-value for bytes. + if response_pb.batch.end_cursor == b'': # Empty-value for bytes. self.next_page_token = None else: - self.next_page_token = base64.urlsafe_b64encode(cursor_as_bytes) + self.next_page_token = base64.urlsafe_b64encode( + response_pb.batch.end_cursor) self._end_cursor = None - if more_results_enum == _NOT_FINISHED: + if response_pb.batch.more_results == _NOT_FINISHED: self._more_results = True - elif more_results_enum in _FINISHED: + elif response_pb.batch.more_results in _FINISHED: self._more_results = False else: raise ValueError('Unexpected value returned for `more_results`.') - return entity_pbs + return [result.entity for result in response_pb.batch.entity_results] def _next_page(self): """Get the next page in the iterator. @@ -492,13 +482,13 @@ def _next_page(self): pb = self._build_protobuf() transaction = self.client.current_transaction - query_results = self.client._connection.run_query( + response_pb = self.client._connection.run_query( query_pb=pb, project=self._query.project, namespace=self._query.namespace, transaction_id=transaction and transaction.id, ) - entity_pbs = self._process_query_results(*query_results) + entity_pbs = self._process_query_results(response_pb) return Page(self, entity_pbs, self._item_to_value) diff --git a/datastore/unit_tests/test__http.py b/datastore/unit_tests/test__http.py index ec4867f7c893..04087c7122c6 100644 --- a/datastore/unit_tests/test__http.py +++ b/datastore/unit_tests/test__http.py @@ -150,7 +150,7 @@ def _make_one(self, client, use_grpc=False): new=use_grpc): return self._get_target_class()(client) - def _verifyProtobufCall(self, called_with, URI, conn): + def _verify_protobuf_call(self, called_with, URI, conn): from google.cloud import _http as connection_module from google.cloud.datastore._http import _CLIENT_INFO @@ -272,7 +272,7 @@ def test_lookup_single_key_empty_response(self): self.assertEqual(len(missing), 0) self.assertEqual(len(deferred), 0) cw = http._called_with - self._verifyProtobufCall(cw, URI, conn) + self._verify_protobuf_call(cw, URI, conn) rq_class = datastore_pb2.LookupRequest request = rq_class() request.ParseFromString(cw['body']) @@ -301,7 +301,7 @@ def test_lookup_single_key_empty_response_w_eventual(self): self.assertEqual(len(missing), 0) self.assertEqual(len(deferred), 0) cw = http._called_with - self._verifyProtobufCall(cw, URI, conn) + self._verify_protobuf_call(cw, URI, conn) rq_class = datastore_pb2.LookupRequest request = rq_class() request.ParseFromString(cw['body']) @@ -342,7 +342,7 @@ def test_lookup_single_key_empty_response_w_transaction(self): self.assertEqual(len(missing), 0) self.assertEqual(len(deferred), 0) cw = http._called_with - self._verifyProtobufCall(cw, URI, conn) + self._verify_protobuf_call(cw, URI, conn) rq_class = datastore_pb2.LookupRequest request = rq_class() request.ParseFromString(cw['body']) @@ -376,7 +376,7 @@ def test_lookup_single_key_nonempty_response(self): self.assertEqual(found.key.path[0].kind, 'Kind') self.assertEqual(found.key.path[0].id, 1234) cw = http._called_with - self._verifyProtobufCall(cw, URI, conn) + self._verify_protobuf_call(cw, URI, conn) rq_class = datastore_pb2.LookupRequest request = rq_class() request.ParseFromString(cw['body']) @@ -405,7 +405,7 @@ def test_lookup_multiple_keys_empty_response(self): self.assertEqual(len(missing), 0) self.assertEqual(len(deferred), 0) cw = http._called_with - self._verifyProtobufCall(cw, URI, conn) + self._verify_protobuf_call(cw, URI, conn) rq_class = datastore_pb2.LookupRequest request = rq_class() request.ParseFromString(cw['body']) @@ -440,7 +440,7 @@ def test_lookup_multiple_keys_w_missing(self): self.assertEqual([missed.key for missed in missing], [key_pb1, key_pb2]) cw = http._called_with - self._verifyProtobufCall(cw, URI, conn) + self._verify_protobuf_call(cw, URI, conn) rq_class = datastore_pb2.LookupRequest request = rq_class() request.ParseFromString(cw['body']) @@ -475,7 +475,7 @@ def test_lookup_multiple_keys_w_deferred(self): self.assertEqual(len(missing), 0) self.assertEqual([def_key for def_key in deferred], [key_pb1, key_pb2]) cw = http._called_with - self._verifyProtobufCall(cw, URI, conn) + self._verify_protobuf_call(cw, URI, conn) self.assertEqual(cw['uri'], URI) self.assertEqual(cw['method'], 'POST') expected_headers = { @@ -497,34 +497,35 @@ def test_run_query_w_eventual_no_transaction(self): from google.cloud.grpc.datastore.v1 import datastore_pb2 from google.cloud.grpc.datastore.v1 import query_pb2 - PROJECT = 'PROJECT' - KIND = 'Nonesuch' - CURSOR = b'\x00' - q_pb = self._make_query_pb(KIND) + project = 'PROJECT' + kind = 'Nonesuch' + cursor = b'\x00' + q_pb = self._make_query_pb(kind) rsp_pb = datastore_pb2.RunQueryResponse() - rsp_pb.batch.end_cursor = CURSOR + rsp_pb.batch.end_cursor = cursor no_more = query_pb2.QueryResultBatch.NO_MORE_RESULTS rsp_pb.batch.more_results = no_more rsp_pb.batch.entity_result_type = query_pb2.EntityResult.FULL + + # Create mock HTTP and client with response. http = Http({'status': '200'}, rsp_pb.SerializeToString()) client = mock.Mock(_http=http, spec=['_http']) + + # Make request. conn = self._make_one(client) - URI = '/'.join([ + response = conn.run_query(project, q_pb, eventual=True) + + # Check the result and verify the callers. + self.assertEqual(response, rsp_pb) + uri = '/'.join([ conn.api_base_url, conn.API_VERSION, 'projects', - PROJECT + ':runQuery', + project + ':runQuery', ]) - pbs, end, more, skipped = conn.run_query(PROJECT, q_pb, - eventual=True) - self.assertEqual(pbs, []) - self.assertEqual(end, CURSOR) - self.assertTrue(more) - self.assertEqual(skipped, 0) cw = http._called_with - self._verifyProtobufCall(cw, URI, conn) - rq_class = datastore_pb2.RunQueryRequest - request = rq_class() + self._verify_protobuf_call(cw, uri, conn) + request = datastore_pb2.RunQueryRequest() request.ParseFromString(cw['body']) self.assertEqual(request.partition_id.namespace_id, '') self.assertEqual(request.query, q_pb) @@ -536,42 +537,44 @@ def test_run_query_wo_eventual_w_transaction(self): from google.cloud.grpc.datastore.v1 import datastore_pb2 from google.cloud.grpc.datastore.v1 import query_pb2 - PROJECT = 'PROJECT' - KIND = 'Nonesuch' - CURSOR = b'\x00' - TRANSACTION = b'TRANSACTION' - q_pb = self._make_query_pb(KIND) + project = 'PROJECT' + kind = 'Nonesuch' + cursor = b'\x00' + transaction = b'TRANSACTION' + q_pb = self._make_query_pb(kind) rsp_pb = datastore_pb2.RunQueryResponse() - rsp_pb.batch.end_cursor = CURSOR + rsp_pb.batch.end_cursor = cursor no_more = query_pb2.QueryResultBatch.NO_MORE_RESULTS rsp_pb.batch.more_results = no_more rsp_pb.batch.entity_result_type = query_pb2.EntityResult.FULL + + # Create mock HTTP and client with response. http = Http({'status': '200'}, rsp_pb.SerializeToString()) client = mock.Mock(_http=http, spec=['_http']) + + # Make request. conn = self._make_one(client) - URI = '/'.join([ + response = conn.run_query( + project, q_pb, transaction_id=transaction) + + # Check the result and verify the callers. + self.assertEqual(response, rsp_pb) + uri = '/'.join([ conn.api_base_url, conn.API_VERSION, 'projects', - PROJECT + ':runQuery', + project + ':runQuery', ]) - pbs, end, more, skipped = conn.run_query( - PROJECT, q_pb, transaction_id=TRANSACTION) - self.assertEqual(pbs, []) - self.assertEqual(end, CURSOR) - self.assertTrue(more) - self.assertEqual(skipped, 0) cw = http._called_with - self._verifyProtobufCall(cw, URI, conn) - rq_class = datastore_pb2.RunQueryRequest - request = rq_class() + self._verify_protobuf_call(cw, uri, conn) + request = datastore_pb2.RunQueryRequest() request.ParseFromString(cw['body']) self.assertEqual(request.partition_id.namespace_id, '') self.assertEqual(request.query, q_pb) self.assertEqual( request.read_options.read_consistency, datastore_pb2.ReadOptions.READ_CONSISTENCY_UNSPECIFIED) - self.assertEqual(request.read_options.transaction, TRANSACTION) + self.assertEqual(request.read_options.transaction, transaction) def test_run_query_w_eventual_and_transaction(self): from google.cloud.grpc.datastore.v1 import datastore_pb2 @@ -595,33 +598,35 @@ def test_run_query_wo_namespace_empty_result(self): from google.cloud.grpc.datastore.v1 import datastore_pb2 from google.cloud.grpc.datastore.v1 import query_pb2 - PROJECT = 'PROJECT' - KIND = 'Nonesuch' - CURSOR = b'\x00' - q_pb = self._make_query_pb(KIND) + project = 'PROJECT' + kind = 'Nonesuch' + cursor = b'\x00' + q_pb = self._make_query_pb(kind) rsp_pb = datastore_pb2.RunQueryResponse() - rsp_pb.batch.end_cursor = CURSOR + rsp_pb.batch.end_cursor = cursor no_more = query_pb2.QueryResultBatch.NO_MORE_RESULTS rsp_pb.batch.more_results = no_more rsp_pb.batch.entity_result_type = query_pb2.EntityResult.FULL + + # Create mock HTTP and client with response. http = Http({'status': '200'}, rsp_pb.SerializeToString()) client = mock.Mock(_http=http, spec=['_http']) + + # Make request. conn = self._make_one(client) - URI = '/'.join([ + response = conn.run_query(project, q_pb) + + # Check the result and verify the callers. + self.assertEqual(response, rsp_pb) + uri = '/'.join([ conn.api_base_url, conn.API_VERSION, 'projects', - PROJECT + ':runQuery', + project + ':runQuery', ]) - pbs, end, more, skipped = conn.run_query(PROJECT, q_pb) - self.assertEqual(pbs, []) - self.assertEqual(end, CURSOR) - self.assertTrue(more) - self.assertEqual(skipped, 0) cw = http._called_with - self._verifyProtobufCall(cw, URI, conn) - rq_class = datastore_pb2.RunQueryRequest - request = rq_class() + self._verify_protobuf_call(cw, uri, conn) + request = datastore_pb2.RunQueryRequest() request.ParseFromString(cw['body']) self.assertEqual(request.partition_id.namespace_id, '') self.assertEqual(request.query, q_pb) @@ -629,32 +634,39 @@ def test_run_query_wo_namespace_empty_result(self): def test_run_query_w_namespace_nonempty_result(self): from google.cloud.grpc.datastore.v1 import datastore_pb2 from google.cloud.grpc.datastore.v1 import entity_pb2 + from google.cloud.grpc.datastore.v1 import query_pb2 - PROJECT = 'PROJECT' - KIND = 'Kind' + project = 'PROJECT' + kind = 'Kind' entity = entity_pb2.Entity() - q_pb = self._make_query_pb(KIND) + q_pb = self._make_query_pb(kind) rsp_pb = datastore_pb2.RunQueryResponse() rsp_pb.batch.entity_results.add(entity=entity) - rsp_pb.batch.entity_result_type = 1 # FULL - rsp_pb.batch.more_results = 3 # NO_MORE_RESULTS + rsp_pb.batch.entity_result_type = query_pb2.EntityResult.FULL + rsp_pb.batch.more_results = query_pb2.QueryResultBatch.NO_MORE_RESULTS + + # Create mock HTTP and client with response. http = Http({'status': '200'}, rsp_pb.SerializeToString()) client = mock.Mock(_http=http, spec=['_http']) + + # Make request. conn = self._make_one(client) - URI = '/'.join([ + namespace = 'NS' + response = conn.run_query(project, q_pb, namespace=namespace) + + # Check the result and verify the callers. + self.assertEqual(response, rsp_pb) + cw = http._called_with + uri = '/'.join([ conn.api_base_url, conn.API_VERSION, 'projects', - PROJECT + ':runQuery', + project + ':runQuery', ]) - pbs = conn.run_query(PROJECT, q_pb, 'NS')[0] - self.assertEqual(len(pbs), 1) - cw = http._called_with - self._verifyProtobufCall(cw, URI, conn) - rq_class = datastore_pb2.RunQueryRequest - request = rq_class() + self._verify_protobuf_call(cw, uri, conn) + request = datastore_pb2.RunQueryRequest() request.ParseFromString(cw['body']) - self.assertEqual(request.partition_id.namespace_id, 'NS') + self.assertEqual(request.partition_id.namespace_id, namespace) self.assertEqual(request.query, q_pb) def test_begin_transaction(self): @@ -675,7 +687,7 @@ def test_begin_transaction(self): ]) self.assertEqual(conn.begin_transaction(PROJECT), TRANSACTION) cw = http._called_with - self._verifyProtobufCall(cw, URI, conn) + self._verify_protobuf_call(cw, URI, conn) rq_class = datastore_pb2.BeginTransactionRequest request = rq_class() request.ParseFromString(cw['body']) @@ -719,7 +731,7 @@ def mock_parse(response): self.assertIs(result, expected_result) cw = http._called_with - self._verifyProtobufCall(cw, URI, conn) + self._verify_protobuf_call(cw, URI, conn) rq_class = datastore_pb2.CommitRequest request = rq_class() request.ParseFromString(cw['body']) @@ -767,7 +779,7 @@ def mock_parse(response): self.assertIs(result, expected_result) cw = http._called_with - self._verifyProtobufCall(cw, URI, conn) + self._verify_protobuf_call(cw, URI, conn) rq_class = datastore_pb2.CommitRequest request = rq_class() request.ParseFromString(cw['body']) @@ -794,7 +806,7 @@ def test_rollback_ok(self): ]) self.assertIsNone(conn.rollback(PROJECT, TRANSACTION)) cw = http._called_with - self._verifyProtobufCall(cw, URI, conn) + self._verify_protobuf_call(cw, URI, conn) rq_class = datastore_pb2.RollbackRequest request = rq_class() request.ParseFromString(cw['body']) @@ -816,7 +828,7 @@ def test_allocate_ids_empty(self): ]) self.assertEqual(conn.allocate_ids(PROJECT, []), []) cw = http._called_with - self._verifyProtobufCall(cw, URI, conn) + self._verify_protobuf_call(cw, URI, conn) rq_class = datastore_pb2.AllocateIdsRequest request = rq_class() request.ParseFromString(cw['body']) @@ -849,7 +861,7 @@ def test_allocate_ids_non_empty(self): self.assertEqual(conn.allocate_ids(PROJECT, before_key_pbs), after_key_pbs) cw = http._called_with - self._verifyProtobufCall(cw, URI, conn) + self._verify_protobuf_call(cw, URI, conn) rq_class = datastore_pb2.AllocateIdsRequest request = rq_class() request.ParseFromString(cw['body']) diff --git a/datastore/unit_tests/test_query.py b/datastore/unit_tests/test_query.py index f3ed774db336..a02e166aba33 100644 --- a/datastore/unit_tests/test_query.py +++ b/datastore/unit_tests/test_query.py @@ -28,13 +28,13 @@ def _get_target_class(): def _make_one(self, *args, **kw): return self._get_target_class()(*args, **kw) - def _makeClient(self, connection=None): + def _make_client(self, connection=None): if connection is None: connection = _Connection() return _Client(self._PROJECT, connection) def test_ctor_defaults(self): - client = self._makeClient() + client = self._make_client() query = self._make_one(client) self.assertIs(query._client, client) self.assertEqual(query.project, client.project) @@ -52,7 +52,7 @@ def test_ctor_explicit(self): _PROJECT = 'OTHER_PROJECT' _KIND = 'KIND' _NAMESPACE = 'OTHER_NAMESPACE' - client = self._makeClient() + client = self._make_client() ancestor = Key('ANCESTOR', 123, project=_PROJECT) FILTERS = [('foo', '=', 'Qux'), ('bar', '<', 17)] PROJECTION = ['foo', 'bar', 'baz'] @@ -81,26 +81,26 @@ def test_ctor_explicit(self): def test_ctor_bad_projection(self): BAD_PROJECTION = object() - self.assertRaises(TypeError, self._make_one, self._makeClient(), + self.assertRaises(TypeError, self._make_one, self._make_client(), projection=BAD_PROJECTION) def test_ctor_bad_order(self): BAD_ORDER = object() - self.assertRaises(TypeError, self._make_one, self._makeClient(), + self.assertRaises(TypeError, self._make_one, self._make_client(), order=BAD_ORDER) def test_ctor_bad_distinct_on(self): BAD_DISTINCT_ON = object() - self.assertRaises(TypeError, self._make_one, self._makeClient(), + self.assertRaises(TypeError, self._make_one, self._make_client(), distinct_on=BAD_DISTINCT_ON) def test_ctor_bad_filters(self): FILTERS_CANT_UNPACK = [('one', 'two')] - self.assertRaises(ValueError, self._make_one, self._makeClient(), + self.assertRaises(ValueError, self._make_one, self._make_client(), filters=FILTERS_CANT_UNPACK) def test_namespace_setter_w_non_string(self): - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) def _assign(val): query.namespace = val @@ -109,12 +109,12 @@ def _assign(val): def test_namespace_setter(self): _NAMESPACE = 'OTHER_NAMESPACE' - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) query.namespace = _NAMESPACE self.assertEqual(query.namespace, _NAMESPACE) def test_kind_setter_w_non_string(self): - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) def _assign(val): query.kind = val @@ -123,21 +123,21 @@ def _assign(val): def test_kind_setter_wo_existing(self): _KIND = 'KIND' - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) query.kind = _KIND self.assertEqual(query.kind, _KIND) def test_kind_setter_w_existing(self): _KIND_BEFORE = 'KIND_BEFORE' _KIND_AFTER = 'KIND_AFTER' - query = self._make_one(self._makeClient(), kind=_KIND_BEFORE) + query = self._make_one(self._make_client(), kind=_KIND_BEFORE) self.assertEqual(query.kind, _KIND_BEFORE) query.kind = _KIND_AFTER self.assertEqual(query.project, self._PROJECT) self.assertEqual(query.kind, _KIND_AFTER) def test_ancestor_setter_w_non_key(self): - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) def _assign(val): query.ancestor = val @@ -150,7 +150,7 @@ def test_ancestor_setter_w_key(self): _NAME = u'NAME' key = Key('KIND', 123, project=self._PROJECT) - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) query.add_filter('name', '=', _NAME) query.ancestor = key self.assertEqual(query.ancestor.path, key.path) @@ -159,22 +159,22 @@ def test_ancestor_deleter_w_key(self): from google.cloud.datastore.key import Key key = Key('KIND', 123, project=self._PROJECT) - query = self._make_one(client=self._makeClient(), ancestor=key) + query = self._make_one(client=self._make_client(), ancestor=key) del query.ancestor self.assertIsNone(query.ancestor) def test_add_filter_setter_w_unknown_operator(self): - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) self.assertRaises(ValueError, query.add_filter, 'firstname', '~~', 'John') def test_add_filter_w_known_operator(self): - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) query.add_filter('firstname', '=', u'John') self.assertEqual(query.filters, [('firstname', '=', u'John')]) def test_add_filter_w_all_operators(self): - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) query.add_filter('leq_prop', '<=', u'val1') query.add_filter('geq_prop', '>=', u'val2') query.add_filter('lt_prop', '<', u'val3') @@ -190,7 +190,7 @@ def test_add_filter_w_all_operators(self): def test_add_filter_w_known_operator_and_entity(self): from google.cloud.datastore.entity import Entity - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) other = Entity() other['firstname'] = u'John' other['lastname'] = u'Smith' @@ -198,7 +198,7 @@ def test_add_filter_w_known_operator_and_entity(self): self.assertEqual(query.filters, [('other', '=', other)]) def test_add_filter_w_whitespace_property_name(self): - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) PROPERTY_NAME = ' property with lots of space ' query.add_filter(PROPERTY_NAME, '=', u'John') self.assertEqual(query.filters, [(PROPERTY_NAME, '=', u'John')]) @@ -206,7 +206,7 @@ def test_add_filter_w_whitespace_property_name(self): def test_add_filter___key__valid_key(self): from google.cloud.datastore.key import Key - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) key = Key('Foo', project=self._PROJECT) query.add_filter('__key__', '=', key) self.assertEqual(query.filters, [('__key__', '=', key)]) @@ -215,47 +215,47 @@ def test_filter___key__not_equal_operator(self): from google.cloud.datastore.key import Key key = Key('Foo', project=self._PROJECT) - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) query.add_filter('__key__', '<', key) self.assertEqual(query.filters, [('__key__', '<', key)]) def test_filter___key__invalid_value(self): - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) self.assertRaises(ValueError, query.add_filter, '__key__', '=', None) def test_projection_setter_empty(self): - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) query.projection = [] self.assertEqual(query.projection, []) def test_projection_setter_string(self): - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) query.projection = 'field1' self.assertEqual(query.projection, ['field1']) def test_projection_setter_non_empty(self): - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) query.projection = ['field1', 'field2'] self.assertEqual(query.projection, ['field1', 'field2']) def test_projection_setter_multiple_calls(self): _PROJECTION1 = ['field1', 'field2'] _PROJECTION2 = ['field3'] - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) query.projection = _PROJECTION1 self.assertEqual(query.projection, _PROJECTION1) query.projection = _PROJECTION2 self.assertEqual(query.projection, _PROJECTION2) def test_keys_only(self): - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) query.keys_only() self.assertEqual(query.projection, ['__key__']) def test_key_filter_defaults(self): from google.cloud.datastore.key import Key - client = self._makeClient() + client = self._make_client() query = self._make_one(client) self.assertEqual(query.filters, []) key = Key('Kind', 1234, project='project') @@ -265,7 +265,7 @@ def test_key_filter_defaults(self): def test_key_filter_explicit(self): from google.cloud.datastore.key import Key - client = self._makeClient() + client = self._make_client() query = self._make_one(client) self.assertEqual(query.filters, []) key = Key('Kind', 1234, project='project') @@ -273,44 +273,44 @@ def test_key_filter_explicit(self): self.assertEqual(query.filters, [('__key__', '>', key)]) def test_order_setter_empty(self): - query = self._make_one(self._makeClient(), order=['foo', '-bar']) + query = self._make_one(self._make_client(), order=['foo', '-bar']) query.order = [] self.assertEqual(query.order, []) def test_order_setter_string(self): - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) query.order = 'field' self.assertEqual(query.order, ['field']) def test_order_setter_single_item_list_desc(self): - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) query.order = ['-field'] self.assertEqual(query.order, ['-field']) def test_order_setter_multiple(self): - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) query.order = ['foo', '-bar'] self.assertEqual(query.order, ['foo', '-bar']) def test_distinct_on_setter_empty(self): - query = self._make_one(self._makeClient(), distinct_on=['foo', 'bar']) + query = self._make_one(self._make_client(), distinct_on=['foo', 'bar']) query.distinct_on = [] self.assertEqual(query.distinct_on, []) def test_distinct_on_setter_string(self): - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) query.distinct_on = 'field1' self.assertEqual(query.distinct_on, ['field1']) def test_distinct_on_setter_non_empty(self): - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) query.distinct_on = ['field1', 'field2'] self.assertEqual(query.distinct_on, ['field1', 'field2']) def test_distinct_on_multiple_calls(self): _DISTINCT_ON1 = ['field1', 'field2'] _DISTINCT_ON2 = ['field3'] - query = self._make_one(self._makeClient()) + query = self._make_one(self._make_client()) query.distinct_on = _DISTINCT_ON1 self.assertEqual(query.distinct_on, _DISTINCT_ON1) query.distinct_on = _DISTINCT_ON2 @@ -320,7 +320,7 @@ def test_fetch_defaults_w_client_attr(self): from google.cloud.datastore.query import Iterator connection = _Connection() - client = self._makeClient(connection) + client = self._make_client(connection) query = self._make_one(client) iterator = query.fetch() @@ -334,8 +334,8 @@ def test_fetch_w_explicit_client(self): from google.cloud.datastore.query import Iterator connection = _Connection() - client = self._makeClient(connection) - other_client = self._makeClient(connection) + client = self._make_client(connection) + other_client = self._make_client(connection) query = self._make_one(client) iterator = query.fetch(limit=7, offset=8, client=other_client) self.assertIsInstance(iterator, Iterator) @@ -443,15 +443,17 @@ def test__process_query_results(self): end_cursor='abcd') self.assertIsNotNone(iterator._end_cursor) - entity_pbs = object() + entity_pbs = [ + _make_entity('Hello', 9998, 'PRAHJEKT'), + ] cursor_as_bytes = b'\x9ai\xe7' cursor = b'mmnn' skipped_results = 4 more_results_enum = query_pb2.QueryResultBatch.NOT_FINISHED - result = iterator._process_query_results( - entity_pbs, cursor_as_bytes, - more_results_enum, skipped_results) - self.assertIs(result, entity_pbs) + response_pb = _make_query_response( + entity_pbs, cursor_as_bytes, more_results_enum, skipped_results) + result = iterator._process_query_results(response_pb) + self.assertEqual(result, entity_pbs) self.assertEqual(iterator._skipped_results, skipped_results) self.assertEqual(iterator.next_page_token, cursor) @@ -464,14 +466,16 @@ def test__process_query_results_done(self): end_cursor='abcd') self.assertIsNotNone(iterator._end_cursor) - entity_pbs = object() + entity_pbs = [ + _make_entity('World', 1234, 'PROJECT'), + ] cursor_as_bytes = b'' skipped_results = 44 more_results_enum = query_pb2.QueryResultBatch.NO_MORE_RESULTS - result = iterator._process_query_results( - entity_pbs, cursor_as_bytes, - more_results_enum, skipped_results) - self.assertIs(result, entity_pbs) + response_pb = _make_query_response( + entity_pbs, cursor_as_bytes, more_results_enum, skipped_results) + result = iterator._process_query_results(response_pb) + self.assertEqual(result, entity_pbs) self.assertEqual(iterator._skipped_results, skipped_results) self.assertIsNone(iterator.next_page_token) @@ -480,9 +484,10 @@ def test__process_query_results_done(self): def test__process_query_results_bad_enum(self): iterator = self._make_one(None, None) more_results_enum = 999 + response_pb = _make_query_response( + [], b'', more_results_enum, 0) with self.assertRaises(ValueError): - iterator._process_query_results( - None, b'', more_results_enum, None) + iterator._process_query_results(response_pb) def test__next_page(self): from google.cloud.iterator import Page @@ -491,7 +496,7 @@ def test__next_page(self): connection = _Connection() more_enum = query_pb2.QueryResultBatch.NOT_FINISHED - result = ([], b'', more_enum, 0) + result = _make_query_response([], b'', more_enum, 0) connection._results = [result] project = 'prujekt' client = _Client(project, connection) @@ -695,3 +700,32 @@ def __init__(self, project, connection, namespace=None): @property def current_transaction(self): pass + + +def _make_entity(kind, id_, project): + from google.cloud.grpc.datastore.v1 import entity_pb2 + + key = entity_pb2.Key() + key.partition_id.project_id = project + elem = key.path.add() + elem.kind = kind + elem.id = id_ + return entity_pb2.Entity(key=key) + + +def _make_query_response( + entity_pbs, cursor_as_bytes, more_results_enum, skipped_results): + from google.cloud.grpc.datastore.v1 import datastore_pb2 + from google.cloud.grpc.datastore.v1 import query_pb2 + + return datastore_pb2.RunQueryResponse( + batch=query_pb2.QueryResultBatch( + skipped_results=skipped_results, + end_cursor=cursor_as_bytes, + more_results=more_results_enum, + entity_results=[ + query_pb2.EntityResult(entity=entity) + for entity in entity_pbs + ], + ), + )