def test__pre_commit_retry_id_already_set_success(self):
from google.cloud.firestore_v1beta1.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_,
options=transaction._client._call_options)
firestore_api.rollback.assert_not_called()
firestore_api.commit.assert_not_called()
def test__options_protobuf_read_only(self):
from google.cloud.firestore_v1beta1.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__options_protobuf_on_retry(self):
from google.cloud.firestore_v1beta1.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)
def test___call__success_second_attempt(self):
from google.api_core import exceptions
from google.cloud.firestore_v1beta1.proto import common_pb2
from google.cloud.firestore_v1beta1.proto import firestore_pb2
from google.cloud.firestore_v1beta1.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])