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 test_commit(self):
# Setup Expected Response
expected_response = {}
expected_response = firestore_pb2.CommitResponse(**expected_response)
# Mock the API response
channel = ChannelStub(responses=[expected_response])
patch = mock.patch("google.api_core.grpc_helpers.create_channel")
with patch as create_channel:
create_channel.return_value = channel
client = firestore_client.FirestoreClient()
# Setup Request
database = "database1789464955"
response = client.commit(database)
assert expected_response == response
assert len(channel.requests) == 1
expected_request = firestore_pb2.CommitRequest(database=database)
actual_request = channel.requests[0][1]
assert expected_request == actual_request
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
