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_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 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, should_recover=_should_recover, should_terminate=_should_terminate, initial_request=initial_request, 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()
def _get_rpc_request(self): if self.resume_token is not None: self._targets["resume_token"] = self.resume_token return firestore_pb2.ListenRequest( database=self._firestore._database_string, add_target=self._targets )