Updating unit tests for GAPIC datastore run_query change.

Author: dhermes

datastore/unit_tests/test__gax.py

@@ -200,46 +200,6 @@ def test_lookup(self):
         self.assertEqual(stub.method_calls,
                          [(request_pb, 'Lookup')])
 
-    def test_run_query(self):
-        return_val = object()
-        stub = _GRPCStub(return_val)
-        datastore_api, _ = self._make_one(stub=stub)
-
-        request_pb = mock.Mock(project_id=None, spec=['project_id'])
-        project = 'PROJECT'
-        result = datastore_api.run_query(project, request_pb)
-        self.assertIs(result, return_val)
-        self.assertEqual(request_pb.project_id, project)
-        self.assertEqual(stub.method_calls,
-                         [(request_pb, 'RunQuery')])
-
-    def _run_query_failure_helper(self, exc, err_class):
-        stub = _GRPCStub(side_effect=exc)
-        datastore_api, _ = self._make_one(stub=stub)
-
-        request_pb = mock.Mock(project_id=None, spec=['project_id'])
-        project = 'PROJECT'
-        with self.assertRaises(err_class):
-            datastore_api.run_query(project, request_pb)
-
-        self.assertEqual(request_pb.project_id, project)
-        self.assertEqual(stub.method_calls,
-                         [(request_pb, 'RunQuery')])
-
-    @unittest.skipUnless(_HAVE_GRPC, 'No gRPC')
-    def test_run_query_invalid_argument(self):
-        from grpc import StatusCode
-        from grpc._channel import _RPCState
-        from google.cloud.exceptions import BadRequest
-        from google.cloud.exceptions import GrpcRendezvous
-
-        details = ('Cannot have inequality filters on multiple '
-                   'properties: [created, priority]')
-        exc_state = _RPCState((), None, None,
-                              StatusCode.INVALID_ARGUMENT, details)
-        exc = GrpcRendezvous(exc_state, None, None, None)
-        self._run_query_failure_helper(exc, BadRequest)
-
 
 @unittest.skipUnless(_HAVE_GRPC, 'No gRPC')
 class TestGAPICDatastoreAPI(unittest.TestCase):
@@ -307,20 +267,13 @@ def test_it(self, make_chan, mock_klass):
 
 class _GRPCStub(object):
 
-    def __init__(self, return_val=None, side_effect=Exception):
+    def __init__(self, return_val=None):
         self.return_val = return_val
-        self.side_effect = side_effect
         self.method_calls = []
 
     def _method(self, request_pb, name):
         self.method_calls.append((request_pb, name))
-        if self.side_effect is Exception:
-            return self.return_val
-        else:
-            raise self.side_effect
+        return self.return_val
 
     def Lookup(self, request_pb):
         return self._method(request_pb, 'Lookup')
-
-    def RunQuery(self, request_pb):
-        return self._method(request_pb, 'RunQuery')

datastore/unit_tests/test__http.py

@@ -161,13 +161,6 @@ def _get_target_class():
 
         return Connection
 
-    def _make_query_pb(self, kind):
-        from google.cloud.proto.datastore.v1 import query_pb2
-
-        pb = query_pb2.Query()
-        pb.kind.add().name = kind
-        return pb
-
     def _make_one(self, client, use_grpc=False):
         with mock.patch('google.cloud.datastore._http._USE_GRPC',
                         new=use_grpc):
@@ -279,14 +272,14 @@ def test_lookup_single_key_empty_response_w_eventual(self):
         self.assertEqual(request.read_options.transaction, b'')
 
     def test_lookup_single_key_empty_response_w_eventual_and_transaction(self):
-        PROJECT = 'PROJECT'
-        TRANSACTION = b'TRANSACTION'
-        key_pb = _make_key_pb(PROJECT)
+        project = 'PROJECT'
+        transaction = b'TRANSACTION'
+        key_pb = _make_key_pb(project)
 
         client = mock.Mock(spec=['_base_url'])
         conn = self._make_one(client)
-        self.assertRaises(ValueError, conn.lookup, PROJECT, key_pb,
-                          eventual=True, transaction_id=TRANSACTION)
+        self.assertRaises(ValueError, conn.lookup, project, [key_pb],
+                          eventual=True, transaction_id=transaction)
 
     def test_lookup_single_key_empty_response_w_transaction(self):
         from google.cloud.proto.datastore.v1 import datastore_pb2
@@ -472,133 +465,152 @@ def test_lookup_multiple_keys_w_deferred(self):
         self.assertEqual(key_pb1, keys[0])
         self.assertEqual(key_pb2, keys[1])
 
+
+class TestHTTPDatastoreAPI(unittest.TestCase):
+
+    @staticmethod
+    def _get_target_class():
+        from google.cloud.datastore._http import HTTPDatastoreAPI
+
+        return HTTPDatastoreAPI
+
+    def _make_one(self, *args, **kwargs):
+        return self._get_target_class()(*args, **kwargs)
+
+    @staticmethod
+    def _make_query_pb(kind):
+        from google.cloud.proto.datastore.v1 import query_pb2
+
+        return query_pb2.Query(
+            kind=[query_pb2.KindExpression(name=kind)],
+        )
+
+    def test_constructor(self):
+        client = object()
+        ds_api = self._make_one(client)
+        self.assertIs(ds_api.client, client)
+
     def test_run_query_w_eventual_no_transaction(self):
         from google.cloud.proto.datastore.v1 import datastore_pb2
+        from google.cloud.proto.datastore.v1 import entity_pb2
         from google.cloud.proto.datastore.v1 import query_pb2
 
         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
-        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
+        query_pb = self._make_query_pb(kind)
+        partition_id = entity_pb2.PartitionId(project_id=project)
+        read_options = datastore_pb2.ReadOptions(
+            read_consistency=datastore_pb2.ReadOptions.EVENTUAL)
+        rsp_pb = datastore_pb2.RunQueryResponse(
+            batch=query_pb2.QueryResultBatch(
+                entity_result_type=query_pb2.EntityResult.FULL,
+                end_cursor=cursor,
+                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, _base_url='test.invalid', spec=['_http', '_base_url'])
 
         # Make request.
-        conn = self._make_one(client)
-        response = conn.run_query(project, q_pb, eventual=True)
+        ds_api = self._make_one(client)
+        response = ds_api.run_query(
+            project, partition_id, read_options, query=query_pb)
 
         # Check the result and verify the callers.
         self.assertEqual(response, rsp_pb)
-        uri = _build_expected_url(conn.api_base_url, project, 'runQuery')
+        uri = _build_expected_url(client._base_url, project, 'runQuery')
         cw = http._called_with
         _verify_protobuf_call(self, cw, uri)
         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.EVENTUAL)
-        self.assertEqual(request.read_options.transaction, b'')
+        self.assertEqual(request.partition_id, partition_id)
+        self.assertEqual(request.query, query_pb)
+        self.assertEqual(request.read_options, read_options)
 
     def test_run_query_wo_eventual_w_transaction(self):
         from google.cloud.proto.datastore.v1 import datastore_pb2
+        from google.cloud.proto.datastore.v1 import entity_pb2
         from google.cloud.proto.datastore.v1 import query_pb2
 
         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
-        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
+        query_pb = self._make_query_pb(kind)
+        partition_id = entity_pb2.PartitionId(project_id=project)
+        read_options = datastore_pb2.ReadOptions(transaction=transaction)
+        rsp_pb = datastore_pb2.RunQueryResponse(
+            batch=query_pb2.QueryResultBatch(
+                entity_result_type=query_pb2.EntityResult.FULL,
+                end_cursor=cursor,
+                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, _base_url='test.invalid', spec=['_http', '_base_url'])
 
         # Make request.
-        conn = self._make_one(client)
-        response = conn.run_query(
-            project, q_pb, transaction_id=transaction)
+        ds_api = self._make_one(client)
+        response = ds_api.run_query(
+            project, partition_id, read_options, query=query_pb)
 
         # Check the result and verify the callers.
         self.assertEqual(response, rsp_pb)
-        uri = _build_expected_url(conn.api_base_url, project, 'runQuery')
+        uri = _build_expected_url(client._base_url, project, 'runQuery')
         cw = http._called_with
         _verify_protobuf_call(self, cw, uri)
         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)
-
-    def test_run_query_w_eventual_and_transaction(self):
-        from google.cloud.proto.datastore.v1 import datastore_pb2
-        from google.cloud.proto.datastore.v1 import query_pb2
-
-        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
-        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
-
-        client = mock.Mock(spec=['_base_url'])
-        conn = self._make_one(client)
-        self.assertRaises(ValueError, conn.run_query, PROJECT, q_pb,
-                          eventual=True, transaction_id=TRANSACTION)
+        self.assertEqual(request.partition_id, partition_id)
+        self.assertEqual(request.query, query_pb)
+        self.assertEqual(request.read_options, read_options)
 
     def test_run_query_wo_namespace_empty_result(self):
         from google.cloud.proto.datastore.v1 import datastore_pb2
+        from google.cloud.proto.datastore.v1 import entity_pb2
         from google.cloud.proto.datastore.v1 import query_pb2
 
         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
-        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
+        query_pb = self._make_query_pb(kind)
+        partition_id = entity_pb2.PartitionId(project_id=project)
+        read_options = datastore_pb2.ReadOptions()
+        rsp_pb = datastore_pb2.RunQueryResponse(
+            batch=query_pb2.QueryResultBatch(
+                entity_result_type=query_pb2.EntityResult.FULL,
+                end_cursor=cursor,
+                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, _base_url='test.invalid', spec=['_http', '_base_url'])
 
         # Make request.
-        conn = self._make_one(client)
-        response = conn.run_query(project, q_pb)
+        ds_api = self._make_one(client)
+        response = ds_api.run_query(
+            project, partition_id, read_options, query=query_pb)
 
         # Check the result and verify the callers.
         self.assertEqual(response, rsp_pb)
-        uri = _build_expected_url(conn.api_base_url, project, 'runQuery')
+        uri = _build_expected_url(client._base_url, project, 'runQuery')
         cw = http._called_with
         _verify_protobuf_call(self, cw, uri)
         request = datastore_pb2.RunQueryRequest()
         request.ParseFromString(cw['body'])
-        self.assertEqual(request.partition_id.namespace_id, '')
-        self.assertEqual(request.query, q_pb)
+        self.assertEqual(request.partition_id, partition_id)
+        self.assertEqual(request.query, query_pb)
+        self.assertEqual(request.read_options, read_options)
 
     def test_run_query_w_namespace_nonempty_result(self):
         from google.cloud.proto.datastore.v1 import datastore_pb2
@@ -607,49 +619,40 @@ def test_run_query_w_namespace_nonempty_result(self):
 
         project = 'PROJECT'
         kind = 'Kind'
-        entity = entity_pb2.Entity()
-        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 = query_pb2.EntityResult.FULL
-        rsp_pb.batch.more_results = query_pb2.QueryResultBatch.NO_MORE_RESULTS
+        namespace = 'NS'
+        query_pb = self._make_query_pb(kind)
+        partition_id = entity_pb2.PartitionId(
+            project_id=project, namespace_id=namespace)
+        read_options = datastore_pb2.ReadOptions()
+        rsp_pb = datastore_pb2.RunQueryResponse(
+            batch=query_pb2.QueryResultBatch(
+                entity_result_type=query_pb2.EntityResult.FULL,
+                entity_results=[
+                    query_pb2.EntityResult(entity=entity_pb2.Entity()),
+                ],
+                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, _base_url='test.invalid', spec=['_http', '_base_url'])
 
         # Make request.
-        conn = self._make_one(client)
-        namespace = 'NS'
-        response = conn.run_query(project, q_pb, namespace=namespace)
+        ds_api = self._make_one(client)
+        response = ds_api.run_query(
+            project, partition_id, read_options, query=query_pb)
 
         # Check the result and verify the callers.
         self.assertEqual(response, rsp_pb)
         cw = http._called_with
-        uri = _build_expected_url(conn.api_base_url, project, 'runQuery')
+        uri = _build_expected_url(client._base_url, project, 'runQuery')
         _verify_protobuf_call(self, cw, uri)
         request = datastore_pb2.RunQueryRequest()
         request.ParseFromString(cw['body'])
-        self.assertEqual(request.partition_id.namespace_id, namespace)
-        self.assertEqual(request.query, q_pb)
-
-
-class TestHTTPDatastoreAPI(unittest.TestCase):
-
-    @staticmethod
-    def _get_target_class():
-        from google.cloud.datastore._http import HTTPDatastoreAPI
-
-        return HTTPDatastoreAPI
-
-    def _make_one(self, *args, **kwargs):
-        return self._get_target_class()(*args, **kwargs)
-
-    def test_constructor(self):
-        client = object()
-        ds_api = self._make_one(client)
-        self.assertIs(ds_api.client, client)
+        self.assertEqual(request.partition_id, partition_id)
+        self.assertEqual(request.query, query_pb)
 
     def test_begin_transaction(self):
         from google.cloud.proto.datastore.v1 import datastore_pb2