async def _create_helper(retry=None, timeout=None):
from google.cloud.firestore_v1 import _helpers
# Create a minimal fake GAPIC with a dummy response.
firestore_api = AsyncMock()
firestore_api.commit.mock_add_spec(spec=["commit"])
firestore_api.commit.return_value = _make_commit_repsonse()
# Attach the fake GAPIC to a real client.
client = _make_client("dignity")
client._firestore_api_internal = firestore_api
# Actually make a document and call create().
document = _make_async_document_reference("foo", "twelve", client=client)
document_data = {"hello": "goodbye", "count": 99}
kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
write_result = await document.create(document_data, **kwargs)
# Verify the response and the mocks.
assert write_result is mock.sentinel.write_result
write_pb = _write_pb_for_create(document._document_path, document_data)
firestore_api.commit.assert_called_once_with(
request={
"database": client._database_string,
"writes": [write_pb],
"transaction": None,
},
metadata=client._rpc_metadata,
**kwargs,
)
async def test_stream_second_response_in_empty_stream(self):
# Create a minimal fake GAPIC with a dummy response.
firestore_api = AsyncMock(spec=["run_query"])
empty_response1 = _make_query_response()
empty_response2 = _make_query_response()
run_query_response = AsyncIter([empty_response1, empty_response2])
firestore_api.run_query.return_value = run_query_response
# Attach the fake GAPIC to a real client.
client = _make_client()
client._firestore_api_internal = firestore_api
# Make a **real** collection reference as parent.
parent = client.collection("dah", "dah", "dum")
query = self._make_one(parent)
get_response = query.stream()
self.assertIsInstance(get_response, types.AsyncGeneratorType)
self.assertEqual([x async for x in get_response], [])
# Verify the mock call.
parent_path, _ = parent._parent_info()
firestore_api.run_query.assert_called_once_with(
request={
"parent": parent_path,
"structured_query": query._to_protobuf(),
"transaction": None,
},
metadata=client._rpc_metadata,
)
async def test_get(self):
# Create a minimal fake GAPIC.
firestore_api = AsyncMock(spec=["run_query"])
# Attach the fake GAPIC to a real client.
client = _make_client()
client._firestore_api_internal = firestore_api
# Make a **real** collection reference as parent.
parent = client.collection("dee")
# Add a dummy response to the minimal fake GAPIC.
_, expected_prefix = parent._parent_info()
name = "{}/sleep".format(expected_prefix)
data = {"snooze": 10}
response_pb = _make_query_response(name=name, data=data)
firestore_api.run_query.return_value = AsyncIter([response_pb])
# Execute the query and check the response.
query = self._make_one(parent)
returned = await query.get()
self.assertIsInstance(returned, list)
self.assertEqual(len(returned), 1)
snapshot = returned[0]
self.assertEqual(snapshot.reference._path, ("dee", "sleep"))
self.assertEqual(snapshot.to_dict(), data)
async def test_failure_first_attempt(self, _sleep):
from google.api_core import exceptions
# Create a minimal fake GAPIC with a dummy result.
firestore_api = AsyncMock()
# Make sure the first request fails with an un-retryable error.
exc = exceptions.ResourceExhausted("We ran out of fries.")
firestore_api.commit.side_effect = exc
# Attach the fake GAPIC to a real client.
client = _make_client("peanut-butter")
client._firestore_api_internal = firestore_api
# Call function and check result.
txn_id = b"\x08\x06\x07\x05\x03\x00\x09-jenny"
with self.assertRaises(exceptions.ResourceExhausted) as exc_info:
await self._call_fut(client, mock.sentinel.write_pbs, txn_id)
self.assertIs(exc_info.exception, exc)
# Verify mocks used.
_sleep.assert_not_called()
firestore_api.commit.assert_called_once_with(
request={
"database": client._database_string,
"writes": mock.sentinel.write_pbs,
"transaction": txn_id,
},
metadata=client._rpc_metadata,
)
async def test_failure_second_attempt(self, _sleep):
from google.api_core import exceptions
# Create a minimal fake GAPIC with a dummy result.
firestore_api = AsyncMock()
# Make sure the first request fails retry-able and second
# fails non-retryable.
exc1 = exceptions.ServiceUnavailable("Come back next time.")
exc2 = exceptions.InternalServerError("Server on fritz.")
firestore_api.commit.side_effect = [exc1, exc2]
# Attach the fake GAPIC to a real client.
client = _make_client("peanut-butter")
client._firestore_api_internal = firestore_api
# Call function and check result.
txn_id = b"the-journey-when-and-where-well-go"
with self.assertRaises(exceptions.InternalServerError) as exc_info:
await self._call_fut(client, mock.sentinel.write_pbs, txn_id)
self.assertIs(exc_info.exception, exc2)
# Verify mocks used.
_sleep.assert_called_once_with(1.0)
# commit() called same way 2 times.
commit_call = mock.call(
request={
"database": client._database_string,
"writes": mock.sentinel.write_pbs,
"transaction": txn_id,
},
metadata=client._rpc_metadata,
)
self.assertEqual(firestore_api.commit.mock_calls, [commit_call, commit_call])
async def test___call__success_first_attempt(self):
to_wrap = AsyncMock(return_value=mock.sentinel.result, spec=[])
wrapped = self._make_one(to_wrap)
txn_id = b"whole-enchilada"
transaction = _make_transaction(txn_id)
result = await 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.
to_wrap.assert_called_once_with(transaction, "a", b="c")
firestore_api = transaction._client._firestore_api
firestore_api.begin_transaction.assert_called_once_with(
request={"database": transaction._client._database_string, "options": None},
metadata=transaction._client._rpc_metadata,
)
firestore_api.rollback.assert_not_called()
firestore_api.commit.assert_called_once_with(
request={
"database": transaction._client._database_string,
"writes": [],
"transaction": txn_id,
},
metadata=transaction._client._rpc_metadata,
)
async def test__begin(self):
from google.cloud.firestore_v1.types import firestore
# Create a minimal fake GAPIC with a dummy result.
firestore_api = AsyncMock()
txn_id = b"to-begin"
response = firestore.BeginTransactionResponse(transaction=txn_id)
firestore_api.begin_transaction.return_value = response
# Attach the fake GAPIC to a real client.
client = _make_client()
client._firestore_api_internal = firestore_api
# Actually make a transaction and ``begin()`` it.
transaction = self._make_one(client)
self.assertIsNone(transaction._id)
ret_val = await transaction._begin()
self.assertIsNone(ret_val)
self.assertEqual(transaction._id, txn_id)
# Verify the called mock.
firestore_api.begin_transaction.assert_called_once_with(
request={"database": client._database_string, "options": None},
metadata=client._rpc_metadata,
)
async def test__pre_commit_retry_id_already_set_success(self):
from google.cloud.firestore_v1.types import common
to_wrap = AsyncMock(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 = await 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.TransactionOptions(
read_write=common.TransactionOptions.ReadWrite(retry_transaction=txn_id1)
)
firestore_api.begin_transaction.assert_called_once_with(
request={
"database": transaction._client._database_string,
"options": options_,
},
metadata=transaction._client._rpc_metadata,
)
firestore_api.rollback.assert_not_called()
firestore_api.commit.assert_not_called()
async def test__pre_commit_failure(self):
exc = RuntimeError("Nope not today.")
to_wrap = AsyncMock(side_effect=exc, spec=[])
wrapped = self._make_one(to_wrap)
txn_id = b"gotta-fail"
transaction = _make_transaction(txn_id)
with self.assertRaises(RuntimeError) as exc_info:
await wrapped._pre_commit(transaction, 10, 20)
self.assertIs(exc_info.exception, exc)
self.assertIsNone(transaction._id)
self.assertEqual(wrapped.current_id, txn_id)
self.assertEqual(wrapped.retry_id, txn_id)
# Verify mocks.
to_wrap.assert_called_once_with(transaction, 10, 20)
firestore_api = transaction._client._firestore_api
firestore_api.begin_transaction.assert_called_once_with(
request={
"database": transaction._client._database_string,
"options": None,
},
metadata=transaction._client._rpc_metadata,
)
firestore_api.rollback.assert_called_once_with(
request={
"database": transaction._client._database_string,
"transaction": txn_id,
},
metadata=transaction._client._rpc_metadata,
)
firestore_api.commit.assert_not_called()
async def _get_w_query_helper(self, retry=None, timeout=None):
from google.cloud.firestore_v1.async_query import AsyncQuery
from google.cloud.firestore_v1 import _helpers
client = AsyncMock(spec=[])
transaction = self._make_one(client)
query = AsyncQuery(parent=AsyncMock(spec=[]))
query.stream = AsyncMock()
kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
result = await transaction.get(query, **kwargs,)
query.stream.assert_called_once_with(
transaction=transaction, **kwargs,
)
self.assertIs(result, query.stream.return_value)
async def test__pre_commit_success(self):
to_wrap = AsyncMock(return_value=mock.sentinel.result, spec=[])
wrapped = self._make_one(to_wrap)
txn_id = b"totes-began"
transaction = _make_transaction(txn_id)
result = await wrapped._pre_commit(transaction, "pos", key="word")
self.assertIs(result, mock.sentinel.result)
self.assertEqual(transaction._id, txn_id)
self.assertEqual(wrapped.current_id, txn_id)
self.assertEqual(wrapped.retry_id, txn_id)
# Verify mocks.
to_wrap.assert_called_once_with(transaction, "pos", key="word")
firestore_api = transaction._client._firestore_api
firestore_api.begin_transaction.assert_called_once_with(
request={
"database": transaction._client._database_string,
"options": None,
},
metadata=transaction._client._rpc_metadata,
)
firestore_api.rollback.assert_not_called()
firestore_api.commit.assert_not_called()
async def test__rollback_failure(self):
from google.api_core import exceptions
# Create a minimal fake GAPIC with a dummy failure.
firestore_api = AsyncMock()
exc = exceptions.InternalServerError("Fire during rollback.")
firestore_api.rollback.side_effect = exc
# Attach the fake GAPIC to a real client.
client = _make_client()
client._firestore_api_internal = firestore_api
# Actually make a transaction and roll it back.
transaction = self._make_one(client)
txn_id = b"roll-bad-server"
transaction._id = txn_id
with self.assertRaises(exceptions.InternalServerError) as exc_info:
await transaction._rollback()
self.assertIs(exc_info.exception, exc)
self.assertIsNone(transaction._id)
self.assertEqual(transaction._write_pbs, [])
# Verify the called mock.
firestore_api.rollback.assert_called_once_with(
request={"database": client._database_string, "transaction": txn_id},
metadata=client._rpc_metadata,
)
async def test_recursive_delete(self):
client = self._make_default_one()
client._firestore_api_internal = AsyncMock(spec=["run_query"])
collection_ref = client.collection("my_collection")
results = []
for index in range(10):
results.append(
RunQueryResponse(document=Document(
name=f"{collection_ref.id}/{index}")))
chunks = [
results[:3],
results[3:6],
results[6:9],
results[9:],
]
def _get_chunk(*args, **kwargs):
return AsyncIter(items=chunks.pop(0))
client._firestore_api_internal.run_query.side_effect = _get_chunk
bulk_writer = mock.MagicMock()
bulk_writer.mock_add_spec(spec=["delete", "close"])
num_deleted = await client.recursive_delete(collection_ref,
bulk_writer=bulk_writer,
chunk_size=3)
self.assertEqual(num_deleted, len(results))
async def _collections_helper(self, retry=None, timeout=None):
from google.cloud.firestore_v1.async_collection import AsyncCollectionReference
from google.cloud.firestore_v1 import _helpers
collection_ids = ["users", "projects"]
class Pager(object):
async def __aiter__(self, **_):
for collection_id in collection_ids:
yield collection_id
firestore_api = AsyncMock()
firestore_api.mock_add_spec(spec=["list_collection_ids"])
firestore_api.list_collection_ids.return_value = Pager()
client = self._make_default_one()
client._firestore_api_internal = firestore_api
kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
collections = [c async for c in client.collections(**kwargs)]
self.assertEqual(len(collections), len(collection_ids))
for collection, collection_id in zip(collections, collection_ids):
self.assertIsInstance(collection, AsyncCollectionReference)
self.assertEqual(collection.parent, None)
self.assertEqual(collection.id, collection_id)
base_path = client._database_string + "/documents"
firestore_api.list_collection_ids.assert_called_once_with(
request={"parent": base_path},
metadata=client._rpc_metadata,
**kwargs,
)
async def _delete_helper(self, **option_kwargs):
from google.cloud.firestore_v1.types import write
# Create a minimal fake GAPIC with a dummy response.
firestore_api = AsyncMock(spec=["commit"])
firestore_api.commit.return_value = self._make_commit_repsonse()
# Attach the fake GAPIC to a real client.
client = _make_client("donut-base")
client._firestore_api_internal = firestore_api
# Actually make a document and call delete().
document = self._make_one("where", "we-are", client=client)
if option_kwargs:
option = client.write_option(**option_kwargs)
delete_time = await document.delete(option=option)
else:
option = None
delete_time = await document.delete()
# Verify the response and the mocks.
self.assertIs(delete_time, mock.sentinel.commit_time)
write_pb = write.Write(delete=document._document_path)
if option is not None:
option.modify_write(write_pb)
firestore_api.commit.assert_called_once_with(
request={
"database": client._database_string,
"writes": [write_pb],
"transaction": None,
},
metadata=client._rpc_metadata,
)
async def _set_helper(self, merge=False, **option_kwargs):
# Create a minimal fake GAPIC with a dummy response.
firestore_api = AsyncMock(spec=["commit"])
firestore_api.commit.return_value = self._make_commit_repsonse()
# Attach the fake GAPIC to a real client.
client = _make_client("db-dee-bee")
client._firestore_api_internal = firestore_api
# Actually make a document and call create().
document = self._make_one("User", "Interface", client=client)
document_data = {"And": 500, "Now": b"\xba\xaa\xaa \xba\xaa\xaa"}
write_result = await document.set(document_data, merge)
# Verify the response and the mocks.
self.assertIs(write_result, mock.sentinel.write_result)
write_pb = self._write_pb_for_set(document._document_path,
document_data, merge)
firestore_api.commit.assert_called_once_with(
request={
"database": client._database_string,
"writes": [write_pb],
"transaction": None,
},
metadata=client._rpc_metadata,
)
async def _set_helper(merge=False, retry=None, timeout=None, **option_kwargs):
from google.cloud.firestore_v1 import _helpers
# Create a minimal fake GAPIC with a dummy response.
firestore_api = AsyncMock(spec=["commit"])
firestore_api.commit.return_value = _make_commit_repsonse()
# Attach the fake GAPIC to a real client.
client = _make_client("db-dee-bee")
client._firestore_api_internal = firestore_api
# Actually make a document and call create().
document = _make_async_document_reference("User",
"Interface",
client=client)
document_data = {"And": 500, "Now": b"\xba\xaa\xaa \xba\xaa\xaa"}
kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
write_result = await document.set(document_data, merge, **kwargs)
# Verify the response and the mocks.
assert write_result is mock.sentinel.write_result
write_pb = _write_pb_for_set(document._document_path, document_data, merge)
firestore_api.commit.assert_called_once_with(
request={
"database": client._database_string,
"writes": [write_pb],
"transaction": None,
},
metadata=client._rpc_metadata,
**kwargs,
)
async def test_get_limit_to_last(self):
from google.cloud import firestore
from google.cloud.firestore_v1.base_query import _enum_from_direction
# Create a minimal fake GAPIC.
firestore_api = AsyncMock(spec=["run_query"])
# Attach the fake GAPIC to a real client.
client = _make_client()
client._firestore_api_internal = firestore_api
# Make a **real** collection reference as parent.
parent = client.collection("dee")
# Add a dummy response to the minimal fake GAPIC.
_, expected_prefix = parent._parent_info()
name = "{}/sleep".format(expected_prefix)
data = {"snooze": 10}
data2 = {"snooze": 20}
response_pb = _make_query_response(name=name, data=data)
response_pb2 = _make_query_response(name=name, data=data2)
firestore_api.run_query.return_value = AsyncIter(
[response_pb2, response_pb])
# Execute the query and check the response.
query = self._make_one(parent)
query = query.order_by(
"snooze",
direction=firestore.AsyncQuery.DESCENDING).limit_to_last(2)
returned = await query.get()
self.assertIsInstance(returned, list)
self.assertEqual(
query._orders[0].direction,
_enum_from_direction(firestore.AsyncQuery.ASCENDING),
)
self.assertEqual(len(returned), 2)
snapshot = returned[0]
self.assertEqual(snapshot.reference._path, ("dee", "sleep"))
self.assertEqual(snapshot.to_dict(), data)
snapshot2 = returned[1]
self.assertEqual(snapshot2.reference._path, ("dee", "sleep"))
self.assertEqual(snapshot2.to_dict(), data2)
# Verify the mock call.
parent_path, _ = parent._parent_info()
firestore_api.run_query.assert_called_once_with(
request={
"parent": parent_path,
"structured_query": query._to_protobuf(),
"transaction": None,
},
metadata=client._rpc_metadata,
)
async def test_get_document_ref(self):
from google.cloud.firestore_v1.async_document import AsyncDocumentReference
client = AsyncMock(spec=["get_all"])
transaction = self._make_one(client)
ref = AsyncDocumentReference("documents", "doc-id")
result = await transaction.get(ref)
client.get_all.assert_called_once_with([ref], transaction=transaction)
self.assertIs(result, client.get_all.return_value)
async def _list_documents_helper(page_size=None, retry=None, timeout=None):
from google.cloud.firestore_v1 import _helpers
from google.api_core.page_iterator_async import AsyncIterator
from google.api_core.page_iterator import Page
from google.cloud.firestore_v1.async_document import AsyncDocumentReference
from google.cloud.firestore_v1.types.document import Document
class _AsyncIterator(AsyncIterator):
def __init__(self, pages):
super(_AsyncIterator, self).__init__(client=None)
self._pages = pages
async def _next_page(self):
if self._pages:
page, self._pages = self._pages[0], self._pages[1:]
return Page(self, page, self.item_to_value)
client = _make_client()
template = client._database_string + "/documents/{}"
document_ids = ["doc-1", "doc-2"]
documents = [
Document(name=template.format(document_id)) for document_id in document_ids
]
iterator = _AsyncIterator(pages=[documents])
firestore_api = AsyncMock()
firestore_api.mock_add_spec(spec=["list_documents"])
firestore_api.list_documents.return_value = iterator
client._firestore_api_internal = firestore_api
collection = _make_async_collection_reference("collection", client=client)
kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
if page_size is not None:
documents = [
i async for i in collection.list_documents(page_size=page_size, **kwargs,)
]
else:
documents = [i async for i in collection.list_documents(**kwargs)]
# Verify the response and the mocks.
assert len(documents) == len(document_ids)
for document, document_id in zip(documents, document_ids):
assert isinstance(document, AsyncDocumentReference)
assert document.parent == collection
assert document.id == document_id
parent, _ = collection._parent_info()
firestore_api.list_documents.assert_called_once_with(
request={
"parent": parent,
"collection_id": collection.id,
"page_size": page_size,
"show_missing": True,
"mask": {"field_paths": None},
},
metadata=client._rpc_metadata,
**kwargs,
)
async def test_add_auto_assigned(self):
from google.cloud.firestore_v1.types import document
from google.cloud.firestore_v1.async_document import AsyncDocumentReference
from google.cloud.firestore_v1 import SERVER_TIMESTAMP
from google.cloud.firestore_v1._helpers import pbs_for_create
# Create a minimal fake GAPIC add attach it to a real client.
firestore_api = AsyncMock(spec=["create_document", "commit"])
write_result = mock.Mock(update_time=mock.sentinel.update_time,
spec=["update_time"])
commit_response = mock.Mock(
write_results=[write_result],
spec=["write_results", "commit_time"],
commit_time=mock.sentinel.commit_time,
)
firestore_api.commit.return_value = commit_response
create_doc_response = document.Document()
firestore_api.create_document.return_value = create_doc_response
client = _make_client()
client._firestore_api_internal = firestore_api
# Actually make a collection.
collection = self._make_one("grand-parent",
"parent",
"child",
client=client)
# Actually call add() on our collection; include a transform to make
# sure transforms during adds work.
document_data = {"been": "here", "now": SERVER_TIMESTAMP}
patch = mock.patch(
"google.cloud.firestore_v1.base_collection._auto_id")
random_doc_id = "DEADBEEF"
with patch as patched:
patched.return_value = random_doc_id
update_time, document_ref = await collection.add(document_data)
# Verify the response and the mocks.
self.assertIs(update_time, mock.sentinel.update_time)
self.assertIsInstance(document_ref, AsyncDocumentReference)
self.assertIs(document_ref._client, client)
expected_path = collection._path + (random_doc_id, )
self.assertEqual(document_ref._path, expected_path)
write_pbs = pbs_for_create(document_ref._document_path, document_data)
firestore_api.commit.assert_called_once_with(
request={
"database": client._database_string,
"writes": write_pbs,
"transaction": None,
},
metadata=client._rpc_metadata,
)
# Since we generate the ID locally, we don't call 'create_document'.
firestore_api.create_document.assert_not_called()
async def test___call__success_second_attempt(self):
from google.api_core import exceptions
from google.cloud.firestore_v1.types import common
from google.cloud.firestore_v1.types import firestore
from google.cloud.firestore_v1.types import write
to_wrap = AsyncMock(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.CommitResponse(write_results=[write.WriteResult()]),
]
# Call the __call__-able ``wrapped``.
result = await 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.TransactionOptions(
read_write=common.TransactionOptions.ReadWrite(retry_transaction=txn_id)
)
self.assertEqual(
firestore_api.begin_transaction.mock_calls,
[
mock.call(
request={"database": db_str, "options": None},
metadata=transaction._client._rpc_metadata,
),
mock.call(
request={"database": db_str, "options": options_},
metadata=transaction._client._rpc_metadata,
),
],
)
firestore_api.rollback.assert_not_called()
commit_call = mock.call(
request={"database": db_str, "writes": [], "transaction": txn_id},
metadata=transaction._client._rpc_metadata,
)
self.assertEqual(firestore_api.commit.mock_calls, [commit_call, commit_call])
async def test___call__failure(self):
from google.api_core import exceptions
from google.cloud.firestore_v1.async_transaction import (
_EXCEED_ATTEMPTS_TEMPLATE,
)
to_wrap = AsyncMock(return_value=mock.sentinel.result, spec=[])
wrapped = self._make_one(to_wrap)
txn_id = b"only-one-shot"
transaction = _make_transaction(txn_id, max_attempts=1)
# Actually force the ``commit`` to fail.
exc = exceptions.Aborted("Contention just once.")
firestore_api = transaction._client._firestore_api
firestore_api.commit.side_effect = exc
# Call the __call__-able ``wrapped``.
with self.assertRaises(ValueError) as exc_info:
await wrapped(transaction, "here", there=1.5)
err_msg = _EXCEED_ATTEMPTS_TEMPLATE.format(transaction._max_attempts)
self.assertEqual(exc_info.exception.args, (err_msg,))
self.assertIsNone(transaction._id)
self.assertEqual(wrapped.current_id, txn_id)
self.assertEqual(wrapped.retry_id, txn_id)
# Verify mocks.
to_wrap.assert_called_once_with(transaction, "here", there=1.5)
firestore_api.begin_transaction.assert_called_once_with(
request={
"database": transaction._client._database_string,
"options": None,
},
metadata=transaction._client._rpc_metadata,
)
firestore_api.rollback.assert_called_once_with(
request={
"database": transaction._client._database_string,
"transaction": txn_id,
},
metadata=transaction._client._rpc_metadata,
)
firestore_api.commit.assert_called_once_with(
request={
"database": transaction._client._database_string,
"writes": [],
"transaction": txn_id,
},
metadata=transaction._client._rpc_metadata,
)
async def _get_w_document_ref_helper(retry=None, timeout=None):
from google.cloud.firestore_v1.async_document import AsyncDocumentReference
from google.cloud.firestore_v1 import _helpers
client = AsyncMock(spec=["get_all"])
transaction = _make_async_transaction(client)
ref = AsyncDocumentReference("documents", "doc-id")
kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
result = await transaction.get(ref, **kwargs)
client.get_all.assert_called_once_with([ref], transaction=transaction, **kwargs)
assert result is client.get_all.return_value
async def _update_helper(self, retry=None, timeout=None, **option_kwargs):
from google.cloud.firestore_v1 import _helpers
from google.cloud.firestore_v1.transforms import DELETE_FIELD
# Create a minimal fake GAPIC with a dummy response.
firestore_api = AsyncMock(spec=["commit"])
firestore_api.commit.return_value = self._make_commit_repsonse()
# Attach the fake GAPIC to a real client.
client = _make_client("potato-chip")
client._firestore_api_internal = firestore_api
# Actually make a document and call create().
document = self._make_one("baked", "Alaska", client=client)
# "Cheat" and use OrderedDict-s so that iteritems() is deterministic.
field_updates = collections.OrderedDict(
(("hello", 1), ("then.do", False), ("goodbye", DELETE_FIELD)))
kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
if option_kwargs:
option = client.write_option(**option_kwargs)
write_result = await document.update(field_updates,
option=option,
**kwargs)
else:
option = None
write_result = await document.update(field_updates, **kwargs)
# Verify the response and the mocks.
self.assertIs(write_result, mock.sentinel.write_result)
update_values = {
"hello": field_updates["hello"],
"then": {
"do": field_updates["then.do"]
},
}
field_paths = list(field_updates.keys())
write_pb = self._write_pb_for_update(document._document_path,
update_values,
sorted(field_paths))
if option is not None:
option.modify_write(write_pb)
firestore_api.commit.assert_called_once_with(
request={
"database": client._database_string,
"writes": [write_pb],
"transaction": None,
},
metadata=client._rpc_metadata,
**kwargs,
)
async def _invoke_get_all(self, client, references, document_pbs, **kwargs):
# Create a minimal fake GAPIC with a dummy response.
firestore_api = AsyncMock(spec=["batch_get_documents"])
response_iterator = AsyncIter(document_pbs)
firestore_api.batch_get_documents.return_value = response_iterator
# Attach the fake GAPIC to a real client.
client._firestore_api_internal = firestore_api
# Actually call get_all().
snapshots = client.get_all(references, **kwargs)
self.assertIsInstance(snapshots, types.AsyncGeneratorType)
return [s async for s in snapshots]
async def _get_all_helper(retry=None, timeout=None):
from google.cloud.firestore_v1 import _helpers
client = AsyncMock(spec=["get_all"])
transaction = _make_async_transaction(client)
ref1, ref2 = mock.Mock(), mock.Mock()
kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
result = await transaction.get_all([ref1, ref2], **kwargs)
client.get_all.assert_called_once_with(
[ref1, ref2], transaction=transaction, **kwargs,
)
assert result is client.get_all.return_value
async def _collections_helper(self, page_size=None):
from google.api_core.page_iterator import Iterator
from google.api_core.page_iterator import Page
from google.cloud.firestore_v1.async_collection import AsyncCollectionReference
# TODO(microgen): https://github.com/googleapis/gapic-generator-python/issues/516
class _Iterator(Iterator):
def __init__(self, pages):
super(_Iterator, self).__init__(client=None)
self._pages = pages
self.collection_ids = pages[0]
def _next_page(self):
if self._pages:
page, self._pages = self._pages[0], self._pages[1:]
return Page(self, page, self.item_to_value)
collection_ids = ["coll-1", "coll-2"]
iterator = _Iterator(pages=[collection_ids])
firestore_api = AsyncMock()
firestore_api.mock_add_spec(spec=["list_collection_ids"])
firestore_api.list_collection_ids.return_value = iterator
client = _make_client()
client._firestore_api_internal = firestore_api
# Actually make a document and call delete().
document = self._make_one("where", "we-are", client=client)
if page_size is not None:
collections = [
c async for c in document.collections(page_size=page_size)
]
else:
collections = [c async for c in document.collections()]
# Verify the response and the mocks.
self.assertEqual(len(collections), len(collection_ids))
for collection, collection_id in zip(collections, collection_ids):
self.assertIsInstance(collection, AsyncCollectionReference)
self.assertEqual(collection.parent, document)
self.assertEqual(collection.id, collection_id)
firestore_api.list_collection_ids.assert_called_once_with(
request={
"parent": document._document_path,
"page_size": page_size
},
metadata=client._rpc_metadata,
)
async def test_empty_update(self):
# Create a minimal fake GAPIC with a dummy response.
firestore_api = AsyncMock(spec=["commit"])
firestore_api.commit.return_value = self._make_commit_repsonse()
# Attach the fake GAPIC to a real client.
client = _make_client("potato-chip")
client._firestore_api_internal = firestore_api
# Actually make a document and call create().
document = self._make_one("baked", "Alaska", client=client)
# "Cheat" and use OrderedDict-s so that iteritems() is deterministic.
field_updates = {}
with self.assertRaises(ValueError):
await document.update(field_updates)
async def _add_helper(self, retry=None, timeout=None):
from google.cloud.firestore_v1.async_document import AsyncDocumentReference
from google.cloud.firestore_v1 import _helpers
# Create a minimal fake GAPIC with a dummy response.
firestore_api = AsyncMock(spec=["commit"])
write_result = mock.Mock(update_time=mock.sentinel.update_time,
spec=["update_time"])
commit_response = mock.Mock(
write_results=[write_result],
spec=["write_results", "commit_time"],
commit_time=mock.sentinel.commit_time,
)
firestore_api.commit.return_value = commit_response
# Attach the fake GAPIC to a real client.
client = _make_client()
client._firestore_api_internal = firestore_api
# Actually make a collection and call add().
collection = self._make_one("parent", client=client)
document_data = {"zorp": 208.75, "i-did-not": b"know that"}
doc_id = "child"
kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
update_time, document_ref = await collection.add(
document_data,
document_id=doc_id,
**kwargs,
)
# Verify the response and the mocks.
self.assertIs(update_time, mock.sentinel.update_time)
self.assertIsInstance(document_ref, AsyncDocumentReference)
self.assertIs(document_ref._client, client)
self.assertEqual(document_ref._path, (collection.id, doc_id))
write_pb = self._write_pb_for_create(document_ref._document_path,
document_data)
firestore_api.commit.assert_called_once_with(
request={
"database": client._database_string,
"writes": [write_pb],
"transaction": None,
},
metadata=client._rpc_metadata,
**kwargs,
)
