fix client_interceptor tests

adap · Sep 27, 2024 · fa3c285 · fa3c285
1 parent bb62bad
commit fa3c285
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 34 deletions.
diff --git a/src/py/flwr/client/app.py b/src/py/flwr/client/app.py
@@ -433,7 +433,7 @@ def _on_backoff(retry_state: RetryState) -> None:
 
                     # Register context for this run
                     node_state.register_context(
-                        run_id=run_id, run=run, flwr_path=flwr_path
+                        run_id=run_id, run=run, flwr_path=flwr_path, fab=fab
                     )
 
                     # Retrieve context for this run

diff --git a/src/py/flwr/client/connection/grpc_bidi/grpc_bidi_connection.py b/src/py/flwr/client/connection/grpc_bidi/grpc_bidi_connection.py
@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 import uuid
+from collections.abc import Iterator
 from logging import DEBUG, ERROR
 from pathlib import Path
 from queue import Queue

diff --git a/src/py/flwr/client/connection/grpc_rere/client_interceptor_test.py b/src/py/flwr/client/connection/grpc_rere/client_interceptor_test.py
@@ -25,7 +25,6 @@
 import grpc
 
 from flwr.common import GRPC_MAX_MESSAGE_LENGTH, serde
-from flwr.common.logger import log
 from flwr.common.message import Message, Metadata
 from flwr.common.record import RecordSet
 from flwr.common.retry_invoker import RetryInvoker, exponential
@@ -166,35 +165,6 @@ def _init_retry_invoker() -> RetryInvoker:
         recoverable_exceptions=grpc.RpcError,
         max_tries=1,
         max_time=None,
-        on_giveup=lambda retry_state: (
-            log(
-                WARN,
-                "Giving up reconnection after %.2f seconds and %s tries.",
-                retry_state.elapsed_time,
-                retry_state.tries,
-            )
-            if retry_state.tries > 1
-            else None
-        ),
-        on_success=lambda retry_state: (
-            log(
-                INFO,
-                "Connection successful after %.2f seconds and %s tries.",
-                retry_state.elapsed_time,
-                retry_state.tries,
-            )
-            if retry_state.tries > 1
-            else None
-        ),
-        on_backoff=lambda retry_state: (
-            log(WARN, "Connection attempt failed, retrying...")
-            if retry_state.tries == 1
-            else log(
-                DEBUG,
-                "Connection attempt failed, retrying in %.2f seconds",
-                retry_state.actual_wait,
-            )
-        ),
     )
 
 
@@ -415,9 +385,8 @@ def test_without_servicer(self) -> None:
             None,
             (self._client_private_key, self._client_public_key),
         ) as conn:
-            _, _, create_node, _, _, _ = conn
-            assert create_node is not None
-            create_node()
+            with self.assertRaises(grpc.RpcError):
+                conn.create_node()
 
             assert self._servicer.received_client_metadata() is None