updates

src/py/flwr/client/client_app.py

@@ -34,10 +34,7 @@ def _inspect_maybe_adapt_client_fn_signature(client_fn: ClientFnExt) -> ClientFn
     client_fn_args = inspect.signature(client_fn).parameters
     first_arg = list(client_fn_args.keys())[0]
 
-    if (
-        len(client_fn_args) != 1
-        or client_fn_args.get(first_arg).annotation is not Context
-    ):
+    if len(client_fn_args) != 1 or client_fn_args[first_arg].annotation is not Context:
         warn_deprecated_feature(
             "`client_fn` now expects a signature `def client_fn(context: Context)`."
             "\nYou provided `client_fn` with signature: "

src/py/flwr/simulation/ray_transport/ray_client_proxy_test.py

@@ -39,7 +39,10 @@
     recordset_to_getpropertiesres,
 )
 from flwr.common.recordset_compat_test import _get_valid_getpropertiesins
-from flwr.simulation.app import _create_node_id_to_partition_mapping
+from flwr.simulation.app import (
+    NodeToPartitionMapping,
+    _create_node_id_to_partition_mapping,
+)
 from flwr.simulation.ray_transport.ray_actor import (
     ClientAppActor,
     VirtualClientEngineActor,
@@ -72,7 +75,9 @@ def get_dummy_client(context: Context) -> Client:
 
 def prep(
     actor_type: Type[VirtualClientEngineActor] = ClientAppActor,
-) -> Tuple[List[RayActorClientProxy], VirtualClientEngineActorPool]:  # pragma: no cover
+) -> Tuple[
+    List[RayActorClientProxy], VirtualClientEngineActorPool, NodeToPartitionMapping
+]:  # pragma: no cover
     """Prepare ClientProxies and pool for tests."""
     client_resources = {"num_cpus": 1, "num_gpus": 0.0}