async def _get_chunk(*args, **kwargs):
start_at = (kwargs["request"]["structured_query"].start_at.
values[0].reference_value)
if collection_1_id in start_at:
return AsyncIter(col_1_chunks.pop(0))
return AsyncIter(col_2_chunks.pop(0))
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_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_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_stream(self, query_class):
query_class.return_value.stream.return_value = AsyncIter(range(3))
collection = self._make_one("collection")
stream_response = collection.stream()
async for _ in stream_response:
pass
query_class.assert_called_once_with(collection)
query_instance = query_class.return_value
query_instance.stream.assert_called_once_with(transaction=None)
async def test_asynccollectionreference_stream_with_transaction(query_class):
query_class.return_value.stream.return_value = AsyncIter(range(3))
collection = _make_async_collection_reference("collection")
transaction = mock.sentinel.txn
stream_response = collection.stream(transaction=transaction)
async for _ in stream_response:
pass
query_class.assert_called_once_with(collection)
query_instance = query_class.return_value
query_instance.stream.assert_called_once_with(transaction=transaction)
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 test_unnecessary_chunkify(self):
client = _make_client()
firestore_api = AsyncMock(spec=["run_query"])
firestore_api.run_query.return_value = AsyncIter([
RunQueryResponse(document=Document(
name=
f"projects/project-project/databases/(default)/documents/asdf/{index}",
), ) for index in range(5)
])
client._firestore_api_internal = firestore_api
query = client.collection("asdf")._query()
async for chunk in query.limit(5)._chunkify(10):
self.assertEqual(len(chunk), 5)
async def _get_partitions_helper(self, retry=None, timeout=None):
from google.cloud.firestore_v1 import _helpers
# Create a minimal fake GAPIC.
firestore_api = AsyncMock(spec=["partition_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("charles")
# Make two **real** document references to use as cursors
document1 = parent.document("one")
document2 = parent.document("two")
# Add cursor pb's to the minimal fake GAPIC.
cursor_pb1 = _make_cursor_pb(([document1], False))
cursor_pb2 = _make_cursor_pb(([document2], False))
firestore_api.partition_query.return_value = AsyncIter(
[cursor_pb1, cursor_pb2])
kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
# Execute the query and check the response.
query = self._make_one(parent)
get_response = query.get_partitions(2, **kwargs)
self.assertIsInstance(get_response, types.AsyncGeneratorType)
returned = [i async for i in get_response]
self.assertEqual(len(returned), 3)
# Verify the mock call.
parent_path, _ = parent._parent_info()
partition_query = self._make_one(
parent,
orders=(query._make_order("__name__", query.ASCENDING), ),
)
firestore_api.partition_query.assert_called_once_with(
request={
"parent": parent_path,
"structured_query": partition_query._to_protobuf(),
"partition_count": 2,
},
metadata=client._rpc_metadata,
**kwargs,
)
async def test_asynccollectionreference_stream_w_retry_timeout(query_class):
from google.api_core.retry import Retry
retry = Retry(predicate=object())
timeout = 123.0
query_class.return_value.stream.return_value = AsyncIter(range(3))
collection = _make_async_collection_reference("collection")
stream_response = collection.stream(retry=retry, timeout=timeout)
async for _ in stream_response:
pass
query_class.assert_called_once_with(collection)
query_instance = query_class.return_value
query_instance.stream.assert_called_once_with(
transaction=None, retry=retry, timeout=timeout,
)
async def _get_helper(self, retry=None, timeout=None):
from google.cloud.firestore_v1 import _helpers
# 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])
kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
# Execute the query and check the response.
query = self._make_one(parent)
returned = await query.get(**kwargs)
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)
# 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,
**kwargs,
)
async def test_stream_w_collection_group(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("charles")
other = client.collection("dora")
# Add two dummy responses to the minimal fake GAPIC.
_, other_prefix = other._parent_info()
name = "{}/bark".format(other_prefix)
data = {"lee": "hoop"}
response_pb1 = _make_query_response(name=name, data=data)
response_pb2 = _make_query_response()
firestore_api.run_query.return_value = AsyncIter(
[response_pb1, response_pb2])
# Execute the query and check the response.
query = self._make_one(parent)
query._all_descendants = True
get_response = query.stream()
self.assertIsInstance(get_response, types.AsyncGeneratorType)
returned = [x async for x in get_response]
self.assertEqual(len(returned), 1)
snapshot = returned[0]
to_match = other.document("bark")
self.assertEqual(snapshot.reference._document_path,
to_match._document_path)
self.assertEqual(snapshot.to_dict(), data)
# 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_stream_with_transaction(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
# Create a real-ish transaction for this client.
transaction = client.transaction()
txn_id = b"\x00\x00\x01-work-\xf2"
transaction._id = txn_id
# Make a **real** collection reference as parent.
parent = client.collection("declaration")
# Add a dummy response to the minimal fake GAPIC.
parent_path, expected_prefix = parent._parent_info()
name = "{}/burger".format(expected_prefix)
data = {"lettuce": b"\xee\x87"}
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)
get_response = query.stream(transaction=transaction)
self.assertIsInstance(get_response, types.AsyncGeneratorType)
returned = [x async for x in get_response]
self.assertEqual(len(returned), 1)
snapshot = returned[0]
self.assertEqual(snapshot.reference._path, ("declaration", "burger"))
self.assertEqual(snapshot.to_dict(), data)
# Verify the mock call.
firestore_api.run_query.assert_called_once_with(
request={
"parent": parent_path,
"structured_query": query._to_protobuf(),
"transaction": txn_id,
},
metadata=client._rpc_metadata,
)
async def test_stream_with_skipped_results(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("talk", "and", "chew-gum")
# Add two dummy responses to the minimal fake GAPIC.
_, expected_prefix = parent._parent_info()
response_pb1 = _make_query_response(skipped_results=1)
name = "{}/clock".format(expected_prefix)
data = {"noon": 12, "nested": {"bird": 10.5}}
response_pb2 = _make_query_response(name=name, data=data)
firestore_api.run_query.return_value = AsyncIter(
[response_pb1, response_pb2])
# Execute the query and check the response.
query = self._make_one(parent)
get_response = query.stream()
self.assertIsInstance(get_response, types.AsyncGeneratorType)
returned = [x async for x in get_response]
self.assertEqual(len(returned), 1)
snapshot = returned[0]
self.assertEqual(snapshot.reference._path,
("talk", "and", "chew-gum", "clock"))
self.assertEqual(snapshot.to_dict(), data)
# 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,
)
def _get_chunk(*args, **kwargs):
return AsyncIter(items=chunks.pop(0))
async def _get_collections(*args, **kwargs):
return AsyncIter([collection_1_id, collection_2_id])
async def _get_helper(
self,
field_paths=None,
use_transaction=False,
not_found=False,
# This should be an impossible case, but we test against it for
# completeness
return_empty=False,
retry=None,
timeout=None,
):
from google.cloud.firestore_v1 import _helpers
from google.cloud.firestore_v1.types import common
from google.cloud.firestore_v1.types import document
from google.cloud.firestore_v1.types import firestore
from google.cloud.firestore_v1.transaction import Transaction
# Create a minimal fake GAPIC with a dummy response.
create_time = 123
update_time = 234
read_time = 345
firestore_api = AsyncMock(spec=["batch_get_documents"])
response = mock.create_autospec(firestore.BatchGetDocumentsResponse)
response.read_time = 345
response.found = mock.create_autospec(document.Document)
response.found.fields = {}
response.found.create_time = create_time
response.found.update_time = update_time
client = _make_client("donut-base")
client._firestore_api_internal = firestore_api
document_reference = self._make_one("where", "we-are", client=client)
response.found.name = None if not_found else document_reference._document_path
response.missing = document_reference._document_path if not_found else None
def WhichOneof(val):
return "missing" if not_found else "found"
response._pb = response
response._pb.WhichOneof = WhichOneof
firestore_api.batch_get_documents.return_value = AsyncIter(
[response] if not return_empty else [])
if use_transaction:
transaction = Transaction(client)
transaction_id = transaction._id = b"asking-me-2"
else:
transaction = None
kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
snapshot = await document_reference.get(
field_paths=field_paths,
transaction=transaction,
**kwargs,
)
self.assertIs(snapshot.reference, document_reference)
if not_found or return_empty:
self.assertIsNone(snapshot._data)
self.assertFalse(snapshot.exists)
self.assertIsNotNone(snapshot.read_time)
self.assertIsNone(snapshot.create_time)
self.assertIsNone(snapshot.update_time)
else:
self.assertEqual(snapshot.to_dict(), {})
self.assertTrue(snapshot.exists)
self.assertIs(snapshot.read_time, read_time)
self.assertIs(snapshot.create_time, create_time)
self.assertIs(snapshot.update_time, update_time)
# Verify the request made to the API
if field_paths is not None:
mask = common.DocumentMask(field_paths=sorted(field_paths))
else:
mask = None
if use_transaction:
expected_transaction_id = transaction_id
else:
expected_transaction_id = None
firestore_api.batch_get_documents.assert_called_once_with(
request={
"database": client._database_string,
"documents": [document_reference._document_path],
"mask": mask,
"transaction": expected_transaction_id,
},
metadata=client._rpc_metadata,
**kwargs,
)
async def _get_chunk(*args, **kwargs):
return AsyncIter(chunks.pop(0))