Add test for get_node_id

adap · Nov 7, 2024 · 070aa1f · 070aa1f
1 parent 1fbf1de
commit 070aa1f
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/src/py/flwr/client/nodestate/in_memory_nodestate.py b/src/py/flwr/client/nodestate/in_memory_nodestate.py
@@ -25,7 +25,7 @@ class InMemoryNodeState(NodeState):
 
     def __init__(self) -> None:
         # Store node_id
-        self.node_id: Optional[int]
+        self.node_id: Optional[int] = None
 
     def set_node_id(self, node_id: Optional[int]) -> None:
         """Set the node ID."""

diff --git a/src/py/flwr/client/nodestate/nodestate_test.py b/src/py/flwr/client/nodestate/nodestate_test.py
@@ -45,6 +45,15 @@ def test_get_set_node_id(self) -> None:
         # Assert
         assert node_id == retrieved_node_id
 
+    def test_get_node_id_fails(self) -> None:
+        """Test get_node_id fails correctly if node_id is not set."""
+        # Prepare
+        state: NodeState = self.state_factory()
+
+        # Execute and assert
+        with self.assertRaises(ValueError):
+            state.get_node_id()
+
 
 class InMemoryStateTest(StateTest):
     """Test InMemoryState implementation."""