def test_listen(self):
# Setup Expected Response
expected_response = {}
expected_response = firestore_pb2.ListenResponse(**expected_response)
# Mock the API response
channel = ChannelStub(responses=[iter([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]")
request = {"database": database}
request = firestore_pb2.ListenRequest(**request)
requests = [request]
response = client.listen(requests)
resources = list(response)
assert len(resources) == 1
assert expected_response == resources[0]
assert len(channel.requests) == 1
actual_requests = channel.requests[0][1]
assert len(actual_requests) == 1
actual_request = list(actual_requests)[0]
assert request == actual_request
def test_listen(self):
# Setup Expected Response
expected_response = {}
expected_response = firestore_pb2.ListenResponse(**expected_response)
# Mock the API response
channel = ChannelStub(responses=[iter([expected_response])])
client = firestore_client.FirestoreClient(channel=channel)
# Setup Request
database = client.database_root_path('[PROJECT]', '[DATABASE]')
request = {'database': database}
request = firestore_pb2.ListenRequest(**request)
requests = [request]
response = client.listen(requests)
resources = list(response)
assert len(resources) == 1
assert expected_response == resources[0]
assert len(channel.requests) == 1
actual_requests = channel.requests[0][1]
assert len(actual_requests) == 1
actual_request = list(actual_requests)[0]
assert request == actual_request
def test_listen_exception(self):
# Mock the API response
channel = ChannelStub(responses=[CustomException()])
client = firestore_client.FirestoreClient(channel=channel)
# Setup request
database = client.database_root_path('[PROJECT]', '[DATABASE]')
request = {'database': database}
request = firestore_pb2.ListenRequest(**request)
requests = [request]
with pytest.raises(CustomException):
client.listen(requests)
def test_listen_exception(self):
# Mock the API response
channel = ChannelStub(responses=[CustomException()])
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]")
request = {"database": database}
request = firestore_pb2.ListenRequest(**request)
requests = [request]
with pytest.raises(CustomException):
client.listen(requests)
def __init__(
self,
document_reference,
firestore,
target,
comparator,
snapshot_callback,
document_snapshot_cls,
document_reference_cls,
BackgroundConsumer=None, # FBO unit testing
ResumableBidiRpc=None, # FBO unit testing
):
"""
Args:
firestore:
target:
comparator:
snapshot_callback: Callback method to process snapshots.
Args:
docs (List(DocumentSnapshot)): A callback that returns the
ordered list of documents stored in this snapshot.
changes (List(str)): A callback that returns the list of
changed documents since the last snapshot delivered for
this watch.
read_time (string): The ISO 8601 time at which this
snapshot was obtained.
document_snapshot_cls: instance of DocumentSnapshot
document_reference_cls: instance of DocumentReference
"""
self._document_reference = document_reference
self._firestore = firestore
self._api = firestore._firestore_api
self._targets = target
self._comparator = comparator
self.DocumentSnapshot = document_snapshot_cls
self.DocumentReference = document_reference_cls
self._snapshot_callback = snapshot_callback
self._closing = threading.Lock()
self._closed = False
def should_recover(exc): # pragma: NO COVER
return (
isinstance(exc, grpc.RpcError)
and exc.code() == grpc.StatusCode.UNAVAILABLE
)
initial_request = firestore_pb2.ListenRequest(
database=self._firestore._database_string, add_target=self._targets
)
if ResumableBidiRpc is None:
ResumableBidiRpc = self.ResumableBidiRpc # FBO unit tests
self._rpc = ResumableBidiRpc(
self._api.transport.listen,
initial_request=initial_request,
should_recover=should_recover,
rpc_metadata=self._firestore._rpc_metadata,
)
self._rpc.add_done_callback(self._on_rpc_done)
# Initialize state for on_snapshot
# The sorted tree of QueryDocumentSnapshots as sent in the last
# snapshot. We only look at the keys.
self.doc_tree = WatchDocTree()
# A map of document names to QueryDocumentSnapshots for the last sent
# snapshot.
self.doc_map = {}
# The accumulates map of document changes (keyed by document name) for
# the current snapshot.
self.change_map = {}
# The current state of the query results.
self.current = False
# We need this to track whether we've pushed an initial set of changes,
# since we should push those even when there are no changes, if there
# aren't docs.
self.has_pushed = False
# The server assigns and updates the resume token.
self.resume_token = None
if BackgroundConsumer is None: # FBO unit tests
BackgroundConsumer = self.BackgroundConsumer
self._consumer = BackgroundConsumer(self._rpc, self.on_snapshot)
self._consumer.start()
