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
def _make_transaction(txn_id, **txn_kwargs):
from google.protobuf import empty_pb2
from google.cloud.firestore_v1beta1.gapic import firestore_client
from google.cloud.firestore_v1beta1.proto import firestore_pb2
from google.cloud.firestore_v1beta1.proto import write_pb2
from google.cloud.firestore_v1beta1.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_commit(self):
from google.cloud.firestore_v1beta1.proto import firestore_pb2
from google.cloud.firestore_v1beta1.proto import write_pb2
# Create a minimal fake GAPIC with a dummy result.
firestore_api = mock.Mock(spec=["commit"])
commit_response = firestore_pb2.CommitResponse(
write_results=[write_pb2.WriteResult(),
write_pb2.WriteResult()])
firestore_api.commit.return_value = commit_response
# Attach the fake GAPIC to a real client.
client = _make_client("grand")
client._firestore_api_internal = firestore_api
# Actually make a batch with some mutations and call commit().
batch = self._make_one(client)
document1 = client.document("a", "b")
batch.create(document1, {"ten": 10, "buck": u"ets"})
document2 = client.document("c", "d", "e", "f")
batch.delete(document2)
write_pbs = batch._write_pbs[::]
write_results = batch.commit()
self.assertEqual(write_results, list(commit_response.write_results))
# Make sure batch has no more "changes".
self.assertEqual(batch._write_pbs, [])
# Verify the mocks.
firestore_api.commit.assert_called_once_with(
client._database_string,
write_pbs,
transaction=None,
metadata=client._rpc_metadata,
)
def test_commit(self):
from google.cloud.firestore_v1beta1.proto import firestore_pb2
from google.cloud.firestore_v1beta1.proto import write_pb2
# Create a minimal fake GAPIC with a dummy result.
firestore_api = mock.Mock(spec=['commit'])
commit_response = firestore_pb2.CommitResponse(write_results=[
write_pb2.WriteResult(),
write_pb2.WriteResult(),
], )
firestore_api.commit.return_value = commit_response
# Attach the fake GAPIC to a real client.
client = _make_client('grand')
client._firestore_api_internal = firestore_api
# Actually make a batch with some mutations and call commit().
batch = self._make_one(client)
document1 = client.document('a', 'b')
batch.create(document1, {'ten': 10, 'buck': u'ets'})
document2 = client.document('c', 'd', 'e', 'f')
batch.delete(document2)
write_pbs = batch._write_pbs[::]
write_results = batch.commit()
self.assertEqual(write_results, list(commit_response.write_results))
# Make sure batch has no more "changes".
self.assertEqual(batch._write_pbs, [])
# Verify the mocks.
firestore_api.commit.assert_called_once_with(
client._database_string,
write_pbs,
transaction=None,
options=client._call_options)
def run_write_test(self, test_proto, desc):
from google.cloud.firestore_v1beta1.proto import firestore_pb2
from google.cloud.firestore_v1beta1.proto import write_pb2
# Create a minimal fake GAPIC with a dummy result.
firestore_api = mock.Mock(spec=['commit'])
commit_response = firestore_pb2.CommitResponse(
write_results=[write_pb2.WriteResult()],
)
firestore_api.commit.return_value = commit_response
kind = test_proto.WhichOneof("test")
call = None
if kind == "create":
tp = test_proto.create
client, doc = self.setup(firestore_api, tp)
data = convert_data(json.loads(tp.json_data))
call = lambda: doc.create(data)
elif kind == "set":
tp = test_proto.set
client, doc = self.setup(firestore_api, tp)
data = convert_data(json.loads(tp.json_data))
# TODO: call doc.set.
elif kind == "update":
tp = test_proto.update
client, doc = self.setup(firestore_api, tp)
data = convert_data(json.loads(tp.json_data))
if tp.HasField("precondition"):
option = convert_precondition(tp.precondition)
else:
option = None
call = lambda: doc.update(data, option)
elif kind == "update_paths":
# Python client doesn't have a way to call update with
# a list of field paths.
pass
else:
assert kind == "delete"
tp = test_proto.delete
client, doc = self.setup(firestore_api, tp)
if tp.HasField("precondition"):
option = convert_precondition(tp.precondition)
else:
option = None
call = lambda: doc.delete(option)
if call is None:
# TODO: remove this when we handle all kinds.
return
if tp.is_error:
# TODO: is there a subclass of Exception we can check for?
with self.assertRaises(Exception):
call()
else:
call()
firestore_api.commit.assert_called_once_with(
client._database_string,
list(tp.request.writes),
transaction=None,
options=client._call_options)
def test_commit(self, mock_create_stub):
# Mock gRPC layer
grpc_stub = mock.Mock()
mock_create_stub.return_value = grpc_stub
client = firestore_client.FirestoreClient()
# Mock request
database = client.database_root_path('[PROJECT]', '[DATABASE]')
writes = []
# Mock response
expected_response = {}
expected_response = firestore_pb2.CommitResponse(**expected_response)
grpc_stub.Commit.return_value = expected_response
response = client.commit(database, writes)
self.assertEqual(expected_response, response)
grpc_stub.Commit.assert_called_once()
args, kwargs = grpc_stub.Commit.call_args
self.assertEqual(len(args), 2)
self.assertEqual(len(kwargs), 1)
self.assertIn('metadata', kwargs)
actual_request = args[0]
expected_request = firestore_pb2.CommitRequest(database=database,
writes=writes)
self.assertEqual(expected_request, actual_request)
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])
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])
client = firestore_client.FirestoreClient(channel=channel)
# Setup Request
database = client.database_root_path('[PROJECT]', '[DATABASE]')
writes = []
response = client.commit(database, writes)
assert expected_response == response
assert len(channel.requests) == 1
expected_request = firestore_pb2.CommitRequest(database=database,
writes=writes)
actual_request = channel.requests[0][1]
assert expected_request == actual_request
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 test__commit(self):
from google.cloud.firestore_v1beta1.gapic import firestore_client
from google.cloud.firestore_v1beta1.proto import firestore_pb2
from google.cloud.firestore_v1beta1.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_as_context_mgr_wo_error(self):
from google.protobuf import timestamp_pb2
from google.cloud.firestore_v1beta1.proto import firestore_pb2
from google.cloud.firestore_v1beta1.proto import write_pb2
firestore_api = mock.Mock(spec=["commit"])
timestamp = timestamp_pb2.Timestamp(seconds=1234567, nanos=123456798)
commit_response = firestore_pb2.CommitResponse(
write_results=[write_pb2.WriteResult(),
write_pb2.WriteResult()],
commit_time=timestamp,
)
firestore_api.commit.return_value = commit_response
client = _make_client()
client._firestore_api_internal = firestore_api
batch = self._make_one(client)
document1 = client.document("a", "b")
document2 = client.document("c", "d", "e", "f")
with batch as ctx_mgr:
self.assertIs(ctx_mgr, batch)
ctx_mgr.create(document1, {"ten": 10, "buck": u"ets"})
ctx_mgr.delete(document2)
write_pbs = batch._write_pbs[::]
self.assertEqual(batch.write_results,
list(commit_response.write_results))
self.assertEqual(batch.commit_time, timestamp)
# Make sure batch has no more "changes".
self.assertEqual(batch._write_pbs, [])
# Verify the mocks.
firestore_api.commit.assert_called_once_with(
client._database_string,
write_pbs,
transaction=None,
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 = client.database_root_path('[PROJECT]', '[DATABASE]')
writes = []
response = client.commit(database, writes)
assert expected_response == response
assert len(channel.requests) == 1
expected_request = firestore_pb2.CommitRequest(database=database,
writes=writes)
actual_request = channel.requests[0][1]
assert expected_request == actual_request
def run_write_test(self, test_proto, desc):
from google.cloud.firestore_v1beta1.proto import firestore_pb2
from google.cloud.firestore_v1beta1.proto import write_pb2
# Create a minimal fake GAPIC with a dummy result.
firestore_api = mock.Mock(spec=['commit'])
commit_response = firestore_pb2.CommitResponse(
write_results=[write_pb2.WriteResult()], )
firestore_api.commit.return_value = commit_response
kind = test_proto.WhichOneof("test")
call = None
if kind == "create":
tp = test_proto.create
client, doc = self.setup(firestore_api, tp)
data = convert_data(json.loads(tp.json_data))
call = functools.partial(doc.create, data)
elif kind == "get":
tp = test_proto.get
client, doc = self.setup(firestore_api, tp)
call = functools.partial(doc.get, None, None)
try:
tp.is_error
except AttributeError:
return
elif kind == "set":
tp = test_proto.set
client, doc = self.setup(firestore_api, tp)
data = convert_data(json.loads(tp.json_data))
if tp.HasField("option"):
merge = True
else:
merge = False
call = functools.partial(doc.set, data, merge)
elif kind == "update":
tp = test_proto.update
client, doc = self.setup(firestore_api, tp)
data = convert_data(json.loads(tp.json_data))
if tp.HasField("precondition"):
option = convert_precondition(tp.precondition)
else:
option = None
call = functools.partial(doc.update, data, option)
elif kind == "update_paths":
# Python client doesn't have a way to call update with
# a list of field paths.
return
else:
assert kind == "delete"
tp = test_proto.delete
client, doc = self.setup(firestore_api, tp)
if tp.HasField("precondition"):
option = convert_precondition(tp.precondition)
else:
option = None
call = functools.partial(doc.delete, option)
if tp.is_error:
# TODO: is there a subclass of Exception we can check for?
with self.assertRaises(Exception):
call()
else:
call()
firestore_api.commit.assert_called_once_with(
client._database_string,
list(tp.request.writes),
transaction=None,
metadata=client._rpc_metadata)
