def test_rollback(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'
transaction = b'-34'
# Mock response
expected_response = datastore_pb2.RollbackResponse()
grpc_stub.Rollback.return_value = expected_response
response = client.rollback(project_id, transaction)
self.assertEqual(expected_response, response)
grpc_stub.Rollback.assert_called_once()
args, kwargs = grpc_stub.Rollback.call_args
self.assertEqual(len(args), 2)
self.assertEqual(len(kwargs), 1)
self.assertIn('metadata', kwargs)
actual_request = args[0]
expected_request = datastore_pb2.RollbackRequest(
project_id=project_id, transaction=transaction)
self.assertEqual(expected_request, actual_request)
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_lookup(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'
keys = []
# Mock response
expected_response = datastore_pb2.LookupResponse()
grpc_stub.Lookup.return_value = expected_response
response = client.lookup(project_id, keys)
self.assertEqual(expected_response, response)
grpc_stub.Lookup.assert_called_once()
args, kwargs = grpc_stub.Lookup.call_args
self.assertEqual(len(args), 2)
self.assertEqual(len(kwargs), 1)
self.assertIn('metadata', kwargs)
actual_request = args[0]
expected_request = datastore_pb2.LookupRequest(project_id=project_id,
keys=keys)
self.assertEqual(expected_request, actual_request)
def test_commit(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'
mode = enums.CommitRequest.Mode.MODE_UNSPECIFIED
mutations = []
# Mock response
index_updates = -1425228195
expected_response = datastore_pb2.CommitResponse(
index_updates=index_updates)
grpc_stub.Commit.return_value = expected_response
response = client.commit(project_id, mode, mutations)
self.assertEqual(expected_response, response)
grpc_stub.Commit.assert_called_once()
args, kwargs = grpc_stub.Commit.call_args
self.assertEqual(len(args), 2)
self.assertEqual(len(kwargs), 1)
self.assertIn('metadata', kwargs)
actual_request = args[0]
expected_request = datastore_pb2.CommitRequest(project_id=project_id,
mode=mode,
mutations=mutations)
self.assertEqual(expected_request, actual_request)
def test_lookup_exception(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'
keys = []
# Mock exception response
grpc_stub.Lookup.side_effect = CustomException()
self.assertRaises(errors.GaxError, client.lookup, project_id, keys)
def test_begin_transaction_exception(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'
# Mock exception response
grpc_stub.BeginTransaction.side_effect = CustomException()
self.assertRaises(errors.GaxError, client.begin_transaction,
project_id)
def test_rollback_exception(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'
transaction = b'-34'
# Mock exception response
grpc_stub.Rollback.side_effect = CustomException()
self.assertRaises(errors.GaxError, client.rollback, project_id,
transaction)
def test_run_query_exception(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 exception response
grpc_stub.RunQuery.side_effect = CustomException()
self.assertRaises(errors.GaxError, client.run_query, project_id,
partition_id)
def test_commit_exception(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'
mode = enums.CommitRequest.Mode.MODE_UNSPECIFIED
mutations = []
# Mock exception response
grpc_stub.Commit.side_effect = CustomException()
self.assertRaises(errors.GaxError, client.commit, project_id, mode,
mutations)
