def test__pre_commit_retry_id_already_set_success(self):
from google.cloud.firestore_v1.proto import common_pb2
to_wrap = mock.Mock(return_value=mock.sentinel.result, spec=[])
wrapped = self._make_one(to_wrap)
txn_id1 = b"already-set"
wrapped.retry_id = txn_id1
txn_id2 = b"ok-here-too"
transaction = _make_transaction(txn_id2)
result = wrapped._pre_commit(transaction)
self.assertIs(result, mock.sentinel.result)
self.assertEqual(transaction._id, txn_id2)
self.assertEqual(wrapped.current_id, txn_id2)
self.assertEqual(wrapped.retry_id, txn_id1)
# Verify mocks.
to_wrap.assert_called_once_with(transaction)
firestore_api = transaction._client._firestore_api
options_ = common_pb2.TransactionOptions(
read_write=common_pb2.TransactionOptions.ReadWrite(
retry_transaction=txn_id1
)
)
firestore_api.begin_transaction.assert_called_once_with(
transaction._client._database_string,
options_=options_,
metadata=transaction._client._rpc_metadata,
)
firestore_api.rollback.assert_not_called()
firestore_api.commit.assert_not_called()
def test__options_protobuf_read_only(self):
from google.cloud.firestore_v1.proto import common_pb2
transaction = self._make_one(mock.sentinel.client, read_only=True)
options_pb = transaction._options_protobuf(None)
expected_pb = common_pb2.TransactionOptions(
read_only=common_pb2.TransactionOptions.ReadOnly()
)
self.assertEqual(options_pb, expected_pb)
def test___call__success_second_attempt(self):
from google.api_core import exceptions
from google.cloud.firestore_v1.proto import common_pb2
from google.cloud.firestore_v1.proto import firestore_pb2
from google.cloud.firestore_v1.proto import write_pb2
to_wrap = mock.Mock(return_value=mock.sentinel.result, spec=[])
wrapped = self._make_one(to_wrap)
txn_id = b"whole-enchilada"
transaction = _make_transaction(txn_id)
# Actually force the ``commit`` to fail on first / succeed on second.
exc = exceptions.Aborted("Contention junction.")
firestore_api = transaction._client._firestore_api
firestore_api.commit.side_effect = [
exc,
firestore_pb2.CommitResponse(write_results=[write_pb2.WriteResult()]),
]
# Call the __call__-able ``wrapped``.
result = wrapped(transaction, "a", b="c")
self.assertIs(result, mock.sentinel.result)
self.assertIsNone(transaction._id)
self.assertEqual(wrapped.current_id, txn_id)
self.assertEqual(wrapped.retry_id, txn_id)
# Verify mocks.
wrapped_call = mock.call(transaction, "a", b="c")
self.assertEqual(to_wrap.mock_calls, [wrapped_call, wrapped_call])
firestore_api = transaction._client._firestore_api
db_str = transaction._client._database_string
options_ = common_pb2.TransactionOptions(
read_write=common_pb2.TransactionOptions.ReadWrite(retry_transaction=txn_id)
)
self.assertEqual(
firestore_api.begin_transaction.mock_calls,
[
mock.call(
db_str, options_=None, metadata=transaction._client._rpc_metadata
),
mock.call(
db_str,
options_=options_,
metadata=transaction._client._rpc_metadata,
),
],
)
firestore_api.rollback.assert_not_called()
commit_call = mock.call(
db_str, [], transaction=txn_id, metadata=transaction._client._rpc_metadata
)
self.assertEqual(firestore_api.commit.mock_calls, [commit_call, commit_call])
def test__options_protobuf_on_retry(self):
from google.cloud.firestore_v1.proto import common_pb2
transaction = self._make_one(mock.sentinel.client)
retry_id = b"hocus-pocus"
options_pb = transaction._options_protobuf(retry_id)
expected_pb = common_pb2.TransactionOptions(
read_write=common_pb2.TransactionOptions.ReadWrite(
retry_transaction=retry_id
)
)
self.assertEqual(options_pb, expected_pb)