async def test_asynctransactional__pre_commit_failure(): exc = RuntimeError("Nope not today.") to_wrap = AsyncMock(side_effect=exc, spec=[]) wrapped = _make_async_transactional(to_wrap) txn_id = b"gotta-fail" transaction = _make_transaction(txn_id) with pytest.raises(RuntimeError) as exc_info: await wrapped._pre_commit(transaction, 10, 20) assert exc_info.value is exc assert transaction._id is None assert wrapped.current_id == txn_id assert 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 _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 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__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__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 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 _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_get_w_query(self): from google.cloud.firestore_v1.async_query import AsyncQuery client = AsyncMock(spec=[]) transaction = self._make_one(client) query = AsyncQuery(parent=AsyncMock(spec=[])) query.stream = AsyncMock() result = await transaction.get(query) query.stream.assert_called_once_with(transaction=transaction) self.assertIs(result, query.stream.return_value)
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 _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 _create_helper(self, 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 = self._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 = self._make_one("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. self.assertIs(write_result, mock.sentinel.write_result) write_pb = self._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_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 _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 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 test_get_all(self): client = AsyncMock(spec=["get_all"]) transaction = self._make_one(client) ref1, ref2 = mock.Mock(), mock.Mock() result = await transaction.get_all([ref1, ref2]) client.get_all.assert_called_once_with([ref1, ref2], transaction=transaction) self.assertIs(result, client.get_all.return_value)
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 _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 _get_w_query_helper(retry=None, timeout=None): from google.cloud.firestore_v1.async_query import AsyncQuery from google.cloud.firestore_v1 import _helpers client = AsyncMock(spec=[]) transaction = _make_async_transaction(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, ) assert result is query.stream.return_value
async def test__commit_with_retry_success_first_attempt(_sleep): from google.cloud.firestore_v1.async_transaction import _commit_with_retry # Create a minimal fake GAPIC with a dummy result. firestore_api = AsyncMock() # Attach the fake GAPIC to a real client. client = _make_client("summer") client._firestore_api_internal = firestore_api # Call function and check result. txn_id = b"cheeeeeez" commit_response = await _commit_with_retry(client, mock.sentinel.write_pbs, txn_id) assert commit_response is firestore_api.commit.return_value # 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_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__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_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 _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__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 _collections_helper(self, page_size=None, retry=None, timeout=None): from google.cloud.firestore_v1 import _helpers from google.cloud.firestore_v1.async_collection import AsyncCollectionReference collection_ids = ["coll-1", "coll-2"] 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 = _make_client() client._firestore_api_internal = firestore_api kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) # 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, **kwargs) ] else: collections = [c async for c in document.collections(**kwargs)] # 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, **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 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__pre_commit_failure_with_rollback_failure(self): from google.api_core import exceptions exc1 = ValueError("I will not be only failure.") to_wrap = AsyncMock(side_effect=exc1, spec=[]) wrapped = self._make_one(to_wrap) txn_id = b"both-will-fail" transaction = _make_transaction(txn_id) # Actually force the ``rollback`` to fail as well. exc2 = exceptions.InternalServerError("Rollback blues.") firestore_api = transaction._client._firestore_api firestore_api.rollback.side_effect = exc2 # Try to ``_pre_commit`` with self.assertRaises(exceptions.InternalServerError) as exc_info: await wrapped._pre_commit(transaction, a="b", c="zebra") self.assertIs(exc_info.exception, exc2) 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="zebra") 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()