[syft/network] add checking NodePeerUpdate type in `NetworkStash.up…

…date`
OpenMined · May 23, 2024 · 0260756 · 0260756
1 parent 479b948
commit 0260756
Showing 1 changed file with 10 additions and 4 deletions.
diff --git a/packages/syft/src/syft/service/network/network_service.py b/packages/syft/src/syft/service/network/network_service.py
@@ -50,6 +50,7 @@
 from ..warnings import CRUDWarning
 from .association_request import AssociationRequestChange
 from .node_peer import NodePeer
+from .node_peer import NodePeerUpdate
 from .routes import HTTPNodeRoute
 from .routes import NodeRoute
 from .routes import NodeRouteType
@@ -87,12 +88,17 @@ def get_by_name(
     def update(
         self,
         credentials: SyftVerifyKey,
-        peer: NodePeer,
+        peer: NodePeer | NodePeerUpdate,
         has_permission: bool = False,
     ) -> Result[NodePeer, str]:
-        valid = self.check_type(peer, NodePeer)
-        if valid.is_err():
-            return Err(SyftError(message=valid.err()))
+        valid_node_peer = self.check_type(peer, NodePeer)
+        valid_node_peer_update = self.check_type(peer, NodePeerUpdate)
+        if valid_node_peer.is_err() and valid_node_peer_update.is_err():
+            return Err(
+                SyftError(
+                    message=f"{type(peer)} does not match required type: NodePeer or NodePeerUpdate"
+                )
+            )
         return super().update(credentials, peer, has_permission=has_permission)
 
     def create_or_update_peer(