From f44e598d6835ee57e410f6dce64d5395565bbae7 Mon Sep 17 00:00:00 2001
From: Danny Hermes <daniel.j.hermes@gmail.com>
Date: Fri, 3 Mar 2017 09:44:35 -0800
Subject: [PATCH] Making datastore Connection.begin_transaction() return
 low-level protobuf.

---
 datastore/google/cloud/datastore/_http.py     |  5 ++--
 .../google/cloud/datastore/transaction.py     |  3 ++-
 datastore/unit_tests/test__http.py            | 26 ++++++++++++-------
 datastore/unit_tests/test_transaction.py      |  5 +++-
 4 files changed, 25 insertions(+), 14 deletions(-)

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