def on_done(self, future: TransferFuture, **_: Any) -> None:
"""On done."""
try:
future.result()
except Exception as exc: # pylint: disable=broad-except
self._on_failure(future, exc)
else:
self._on_success(future)
def setUp(self):
self.subscriber = RecordingSubscriber()
self.second_subscriber = RecordingSubscriber()
self.call_args = CallArgs(
subscribers=[self.subscriber, self.second_subscriber])
self.transfer_meta = TransferMeta(self.call_args)
self.transfer_future = TransferFuture(self.transfer_meta)
def _get_transfer_future(self, **kwargs):
components = {
'meta': self.meta,
'coordinator': self.coordinator,
}
for component_name, component in kwargs.items():
components[component_name] = component
return TransferFuture(**components)
def _get_future_with_components(self, call_args):
transfer_id = self._id_counter
# Creates a new transfer future along with its components
transfer_coordinator = TransferCoordinator(transfer_id=transfer_id)
# Track the transfer coordinator for transfers to manage.
self._coordinator_controller.add_transfer_coordinator(
transfer_coordinator)
# Also make sure that the transfer coordinator is removed once
# the transfer completes so it does not stick around in memory.
transfer_coordinator.add_done_callback(
self._coordinator_controller.remove_transfer_coordinator,
transfer_coordinator)
components = {
'meta': TransferMeta(call_args, transfer_id=transfer_id),
'coordinator': transfer_coordinator
}
transfer_future = TransferFuture(**components)
return transfer_future, components
def get_transfer_future(self, call_args=None):
return TransferFuture(meta=TransferMeta(call_args),
coordinator=self.transfer_coordinator)
def _on_success(self, future: TransferFuture) -> None:
try:
self._delete_source(future)
except Exception as exc: # pylint: disable=broad-except
future.set_exception(exc)
