def test_run_query(self, mock_create_stub): # Mock gRPC layer grpc_stub = mock.Mock() mock_create_stub.return_value = grpc_stub client = datastore_client.DatastoreClient() # Mock request project_id = 'projectId-1969970175' partition_id = entity_pb2.PartitionId() # Mock response expected_response = datastore_pb2.RunQueryResponse() grpc_stub.RunQuery.return_value = expected_response response = client.run_query(project_id, partition_id) self.assertEqual(expected_response, response) grpc_stub.RunQuery.assert_called_once() args, kwargs = grpc_stub.RunQuery.call_args self.assertEqual(len(args), 2) self.assertEqual(len(kwargs), 1) self.assertIn('metadata', kwargs) actual_request = args[0] expected_request = datastore_pb2.RunQueryRequest( project_id=project_id, partition_id=partition_id) self.assertEqual(expected_request, actual_request)
def test_run_query_w_namespace_nonempty_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 = '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 # 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) # 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') _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)
def create_response(entities, end_cursor, finish): """Creates a query response for a given batch of scatter entities.""" resp = datastore_pb2.RunQueryResponse() if finish: resp.batch.more_results = query_pb2.QueryResultBatch.NO_MORE_RESULTS else: resp.batch.more_results = query_pb2.QueryResultBatch.NOT_FINISHED resp.batch.end_cursor = end_cursor for entity_result in entities: resp.batch.entity_results.add().CopyFrom(entity_result) return resp
def _make_query_response(entity_pbs, cursor_as_bytes, more_results_enum, skipped_results): from google.cloud.proto.datastore.v1 import datastore_pb2 from google.cloud.proto.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 ], ), )
def test_run_query_w_namespace_nonempty_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 = 'Kind' 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 = _make_requests_session( [_make_response(content=rsp_pb.SerializeToString())]) client = mock.Mock(_http=http, _base_url='test.invalid', spec=['_http', '_base_url']) # Make request. 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(client._base_url, project, 'runQuery') request = _verify_protobuf_call(http, uri, datastore_pb2.RunQueryRequest()) self.assertEqual(request.partition_id, partition_id) self.assertEqual(request.query, query_pb)
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' 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. 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(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, 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 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 # 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) 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', ]) cw = http._called_with 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)
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 conn = self._make_one(object()) self.assertRaises(ValueError, conn.run_query, PROJECT, q_pb, eventual=True, transaction_id=TRANSACTION)
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 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 # 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) # Check the result and verify the callers. self.assertEqual(response, rsp_pb) uri = _build_expected_url(conn.api_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'')
def make_stats_response(self, property_map): resp = datastore_pb2.RunQueryResponse() entity_result = resp.batch.entity_results.add() datastore_helper.add_properties(entity_result.entity, property_map) return resp
def transient_retriable_datastore_failure(self, req): if self._transient_fail_count: self._transient_fail_count -= 1 raise random.choice(self._retriable_errors) else: return datastore_pb2.RunQueryResponse()
def transient_retriable_datastore_failure(self, req): if self._transient_fail_count: self._transient_fail_count -= 1 raise RPCError("dummy", code_pb2.INTERNAL, "failed") else: return datastore_pb2.RunQueryResponse()
def transient_datastore_failure(self, req): if self._transient_fail_count: self._transient_fail_count -= 1 raise RPCError("dummy", 500, "failed") else: return datastore_pb2.RunQueryResponse()