diff --git a/datastore/google/cloud/datastore/_http.py b/datastore/google/cloud/datastore/_http.py index 4c384ad4d4ef..2650321068bf 100644 --- a/datastore/google/cloud/datastore/_http.py +++ b/datastore/google/cloud/datastore/_http.py @@ -399,12 +399,11 @@ def begin_transaction(self, project): :type project: str :param project: The project to which the transaction applies. - :rtype: bytes + :rtype: :class:`.datastore_pb2.BeginTransactionResponse` :returns: The serialized transaction that was begun. """ request = _datastore_pb2.BeginTransactionRequest() - response = self._datastore_api.begin_transaction(project, request) - return response.transaction + return self._datastore_api.begin_transaction(project, request) def commit(self, project, request, transaction_id): """Commit mutations in context of current transaction (if any). diff --git a/datastore/google/cloud/datastore/transaction.py b/datastore/google/cloud/datastore/transaction.py index 48a044ba789c..c0781a733711 100644 --- a/datastore/google/cloud/datastore/transaction.py +++ b/datastore/google/cloud/datastore/transaction.py @@ -184,8 +184,9 @@ def begin(self): """ super(Transaction, self).begin() try: - self._id = self._client._connection.begin_transaction( + response_pb = self._client._connection.begin_transaction( self.project) + self._id = response_pb.transaction except: # noqa: E722 do not use bare except, specify exception instead self._status = self._ABORTED raise diff --git a/datastore/unit_tests/test__http.py b/datastore/unit_tests/test__http.py index fffc9c2bdb0b..0103e4883479 100644 --- a/datastore/unit_tests/test__http.py +++ b/datastore/unit_tests/test__http.py @@ -672,25 +672,33 @@ def test_run_query_w_namespace_nonempty_result(self): def test_begin_transaction(self): from google.cloud.proto.datastore.v1 import datastore_pb2 - PROJECT = 'PROJECT' - TRANSACTION = b'TRANSACTION' + project = 'PROJECT' + transaction = b'TRANSACTION' rsp_pb = datastore_pb2.BeginTransactionResponse() - rsp_pb.transaction = TRANSACTION + rsp_pb.transaction = transaction + + # 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.begin_transaction(project) + + # Check the result and verify the callers. + self.assertEqual(response, rsp_pb) + uri = '/'.join([ conn.api_base_url, conn.API_VERSION, 'projects', - PROJECT + ':beginTransaction', + project + ':beginTransaction', ]) - self.assertEqual(conn.begin_transaction(PROJECT), TRANSACTION) cw = http._called_with - self._verify_protobuf_call(cw, URI, conn) - rq_class = datastore_pb2.BeginTransactionRequest - request = rq_class() + self._verify_protobuf_call(cw, uri, conn) + request = datastore_pb2.BeginTransactionRequest() request.ParseFromString(cw['body']) + # The RPC-over-HTTP request does not set the project in the request. + self.assertEqual(request.project_id, u'') def test_commit_wo_transaction(self): from google.cloud.proto.datastore.v1 import datastore_pb2 diff --git a/datastore/unit_tests/test_transaction.py b/datastore/unit_tests/test_transaction.py index a4d278a421c3..db7bf6bcc97f 100644 --- a/datastore/unit_tests/test_transaction.py +++ b/datastore/unit_tests/test_transaction.py @@ -206,9 +206,12 @@ def __init__(self, xact_id=123, keys=()): mutation_results=mutation_results) def begin_transaction(self, project): + import mock + self._begun = project if self._side_effect is None: - return self._xact_id + return mock.Mock( + transaction=self._xact_id, spec=['transaction']) else: raise self._side_effect