fixed tests

google/cloud/datastore/query.py

@@ -784,8 +784,9 @@ def _next_page(self):
             transaction_id = None
         else:
             transaction_id = transaction.id
-            # if transaction hasn't been initialized, initialize it as part of this request
-            new_transaction_options = transaction._options
+            if transaction._begin_later and transaction._status == transaction._INITIAL:
+                # if transaction hasn't been initialized, initialize it as part of this request
+                new_transaction_options = transaction._options
         read_options = helpers.get_read_options(
             self._eventual, transaction_id, self._read_time, new_transaction_options
         )

tests/unit/test_aggregation.py

@@ -471,7 +471,9 @@ def _next_page_helper(txn_id=None, retry=None, timeout=None, database_id=None):
     if txn_id is None:
         client = _Client(project, datastore_api=ds_api, database=database_id)
     else:
-        transaction = mock.Mock(id=txn_id, spec=["id"])
+        transaction = mock.Mock(
+            id=txn_id, _begin_later=False, spec=["id", "_begin_later"]
+        )
         client = _Client(
             project, datastore_api=ds_api, transaction=transaction, database=database_id
         )

tests/unit/test_query.py

@@ -667,7 +667,7 @@ def test_eventual_transaction_fails(database_id):
 @pytest.mark.parametrize("database_id", [None, "somedb"])
 def test_transaction_id_populated(database_id):
     """
-    When an aggregation is run in the context of a transaction, the transaction
+    When an query is run in the context of a transaction, the transaction
     ID should be populated in the request.
     """
     import mock
@@ -699,7 +699,7 @@ def test_transaction_id_populated(database_id):
 
 
 @pytest.mark.parametrize("database_id", [None, "somedb"])
-def test_transaction_begin_later(database_id):
+def test_query_transaction_begin_later(database_id):
     """
     When an aggregation is run in the context of a transaction with begin_later=True,
     the new_transaction field should be populated in the request read_options.
@@ -711,7 +711,7 @@ def test_transaction_begin_later(database_id):
     transaction = mock.Mock()
     transaction.id = None
     transaction._begin_later = True
-    transaction._state = transaction._INITIAL
+    transaction._status = transaction._INITIAL
     transaction._options = TransactionOptions(read_only=TransactionOptions.ReadOnly())
 
     mock_datastore_api = mock.Mock()
@@ -926,7 +926,9 @@ def _next_page_helper(
     if txn_id is None:
         client = _Client(project, database=database, datastore_api=ds_api)
     else:
-        transaction = mock.Mock(id=txn_id, spec=["id"])
+        transaction = mock.Mock(
+            id=txn_id, _begin_later=False, spec=["id", "_begin_later"]
+        )
         client = _Client(
             project, database=database, datastore_api=ds_api, transaction=transaction
         )

tests/unit/test_transaction.py

@@ -491,70 +491,6 @@ def test_transaction_put_read_only(database_id):
         xact.put(entity)
 
 
-@pytest.mark.parametrize("database_id", [None, "somedb"])
-def test_transaction_begin_later_with_delete(database_id):
-    """
-    begin_later transactions should begin on delete
-    """
-    project = "PROJECT"
-    id_ = 912830
-    ds_api = _make_datastore_api(xact_id=id_)
-    client = _Client(project, datastore_api=ds_api, database=database_id)
-    xact = _make_transaction(client, begin_later=True)
-
-    fake_key = mock.Mock()
-
-    with mock.patch("google.cloud.datastore.batch.Batch.delete") as delete:
-        with xact:
-            assert xact._status == xact._INITIAL
-            assert xact.id is None
-            xact.delete(fake_key)
-            # call should have started transaction
-            assert xact._status == xact._IN_PROGRESS
-            assert xact.id == id_
-            # super class delete should have been called
-            assert delete.call_count == 1
-            assert delete.call_args == mock.call(fake_key)
-    # should be finalized after context manager block
-    assert xact._status == xact._FINISHED
-    assert xact.id is None
-    # should have committed
-    assert ds_api.commit.call_count == 1
-    assert ds_api.rollback.call_count == 0
-
-
-@pytest.mark.parametrize("database_id", [None, "somedb"])
-def test_transaction_begin_later_with_put(database_id):
-    """
-    begin_later transactions should begin on put
-    """
-    project = "PROJECT"
-    id_ = 912830
-    ds_api = _make_datastore_api(xact_id=id_)
-    client = _Client(project, datastore_api=ds_api, database=database_id)
-    xact = _make_transaction(client, begin_later=True)
-
-    fake_entity = mock.Mock()
-
-    with mock.patch("google.cloud.datastore.batch.Batch.put") as put:
-        with xact:
-            assert xact._status == xact._INITIAL
-            assert xact.id is None
-            xact.put(fake_entity)
-            # call should have started transaction
-            assert xact._status == xact._IN_PROGRESS
-            assert xact.id == id_
-            # super class put should have been called
-            assert put.call_count == 1
-            assert put.call_args == mock.call(fake_entity)
-    # should be finalized after context manager block
-    assert xact._status == xact._FINISHED
-    assert xact.id is None
-    # should have committed
-    assert ds_api.commit.call_count == 1
-    assert ds_api.rollback.call_count == 0
-
-
 def _make_key(kind, id_, project, database=None):
     from google.cloud.datastore_v1.types import entity as entity_pb2