def _make_transaction(txn_id, **txn_kwargs): from google.protobuf import empty_pb2 from google.cloud.firestore_v1.gapic import firestore_client from google.cloud.firestore_v1.proto import firestore_pb2 from google.cloud.firestore_v1.proto import write_pb2 from google.cloud.firestore_v1.transaction import Transaction # Create a fake GAPIC ... firestore_api = mock.create_autospec( firestore_client.FirestoreClient, instance=True ) # ... with a dummy ``BeginTransactionResponse`` result ... begin_response = firestore_pb2.BeginTransactionResponse(transaction=txn_id) firestore_api.begin_transaction.return_value = begin_response # ... and a dummy ``Rollback`` result ... firestore_api.rollback.return_value = empty_pb2.Empty() # ... and a dummy ``Commit`` result. commit_response = firestore_pb2.CommitResponse( write_results=[write_pb2.WriteResult()] ) firestore_api.commit.return_value = commit_response # Attach the fake GAPIC to a real client. client = _make_client() client._firestore_api_internal = firestore_api return Transaction(client, **txn_kwargs)
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__commit(self): from google.cloud.firestore_v1.gapic import firestore_client from google.cloud.firestore_v1.proto import firestore_pb2 from google.cloud.firestore_v1.proto import write_pb2 # Create a minimal fake GAPIC with a dummy result. firestore_api = mock.create_autospec( firestore_client.FirestoreClient, instance=True ) commit_response = firestore_pb2.CommitResponse( write_results=[write_pb2.WriteResult()] ) firestore_api.commit.return_value = commit_response # Attach the fake GAPIC to a real client. client = _make_client("phone-joe") client._firestore_api_internal = firestore_api # Actually make a transaction with some mutations and call _commit(). transaction = self._make_one(client) txn_id = b"under-over-thru-woods" transaction._id = txn_id document = client.document("zap", "galaxy", "ship", "space") transaction.set(document, {"apple": 4.5}) write_pbs = transaction._write_pbs[::] write_results = transaction._commit() self.assertEqual(write_results, list(commit_response.write_results)) # Make sure transaction has no more "changes". self.assertIsNone(transaction._id) self.assertEqual(transaction._write_pbs, []) # Verify the mocks. firestore_api.commit.assert_called_once_with( client._database_string, write_pbs, transaction=txn_id, metadata=client._rpc_metadata, )
def _mock_firestore_api(): firestore_api = mock.Mock(spec=["commit"]) commit_response = firestore_pb2.CommitResponse( write_results=[write_pb2.WriteResult()]) firestore_api.commit.return_value = commit_response return firestore_api