From fa3c28532ce1ed75e66811131ba67c0be70bd751 Mon Sep 17 00:00:00 2001
From: Heng Pan <pan@flower.ai>
Date: Thu, 26 Sep 2024 14:24:09 +0100
Subject: [PATCH] fix client_interceptor tests

---
 src/py/flwr/client/app.py                     |  2 +-
 .../grpc_bidi/grpc_bidi_connection.py         |  1 +
 .../grpc_rere/client_interceptor_test.py      | 35 ++-----------------
 3 files changed, 4 insertions(+), 34 deletions(-)

diff --git a/src/py/flwr/client/app.py b/src/py/flwr/client/app.py
index 328edabd7723..924f52ca70c3 100644
--- 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
index d2591c482dc6..ec0707fd083a 100644
--- 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
index c400f9fa4ff8..1cafe4872832 100644
--- 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