From 3df4ec3927e28b65336593e7bae3ec749649fdf5 Mon Sep 17 00:00:00 2001
From: Jonathan Zernik <yzernik@gmail.com>
Date: Mon, 20 Sep 2021 15:13:29 -0700
Subject: [PATCH] Got itest working for get peer by address (#1364)

---
 itests/tests/test_squeak_node.py                 | 11 +++++++++++
 proto/squeak_admin.proto                         | 14 ++++++++++++++
 squeaknode/admin/squeak_admin_server_handler.py  | 14 ++++++++++++++
 squeaknode/admin/squeak_admin_server_servicer.py |  3 +++
 squeaknode/node/squeak_controller.py             |  3 +++
 5 files changed, 45 insertions(+)

diff --git a/itests/tests/test_squeak_node.py b/itests/tests/test_squeak_node.py
index 82093437..24c67dcf 100644
--- a/itests/tests/test_squeak_node.py
+++ b/itests/tests/test_squeak_node.py
@@ -497,6 +497,17 @@ def test_create_peer(admin_stub):
     ]
     assert "fake_host" in peer_hosts
 
+    # Get the new peer by address
+    get_peer_by_address_response = admin_stub.GetPeerByAddress(
+        squeak_admin_pb2.GetPeerByAddressRequest(
+            peer_address=squeak_admin_pb2.PeerAddress(
+                host="fake_host",
+                port=1234,
+            )
+        )
+    )
+    assert get_peer_by_address_response.squeak_peer.peer_name == "fake_peer_name"
+
 
 def test_create_peer_empty_name(admin_stub):
     # Try to create a new signing profile with an empty name
diff --git a/proto/squeak_admin.proto b/proto/squeak_admin.proto
index 79188659..c428438b 100644
--- a/proto/squeak_admin.proto
+++ b/proto/squeak_admin.proto
@@ -172,6 +172,10 @@ service SqueakAdmin {
   */
   rpc GetPeer (GetPeerRequest) returns (GetPeerReply) {}
 
+  /** sqkadmin: `getpeerbyaddress`
+  */
+  rpc GetPeerByAddress (GetPeerByAddressRequest) returns (GetPeerByAddressReply) {}
+
   /** sqkadmin: `getpeers`
   */
   rpc GetPeers (GetPeersRequest) returns (GetPeersReply) {}
@@ -695,6 +699,16 @@ message GetPeerReply {
     SqueakPeer squeak_peer = 1;
 }
 
+message GetPeerByAddressRequest {
+    /// Peer address
+    PeerAddress peer_address = 1;
+}
+
+message GetPeerByAddressReply {
+    /// The peer
+    SqueakPeer squeak_peer = 1;
+}
+
 message GetPeersRequest {
 }
 
diff --git a/squeaknode/admin/squeak_admin_server_handler.py b/squeaknode/admin/squeak_admin_server_handler.py
index 2b171cac..b2800569 100644
--- a/squeaknode/admin/squeak_admin_server_handler.py
+++ b/squeaknode/admin/squeak_admin_server_handler.py
@@ -531,6 +531,20 @@ class SqueakAdminServerHandler(object):
             squeak_peer=squeak_peer_msg,
         )
 
+    def handle_get_squeak_peer_by_address(self, request):
+        peer_address = message_to_peer_address(request.peer_address)
+        logger.info(
+            "Handle get squeak peer with address: {}".format(peer_address))
+        squeak_peer = self.squeak_controller.get_peer_by_address(peer_address)
+        if squeak_peer is None:
+            return squeak_admin_pb2.GetPeerReply(
+                squeak_peer=None,
+            )
+        squeak_peer_msg = squeak_peer_to_message(squeak_peer)
+        return squeak_admin_pb2.GetPeerByAddressReply(
+            squeak_peer=squeak_peer_msg,
+        )
+
     def handle_get_squeak_peers(self, request):
         logger.info("Handle get squeak peers")
         squeak_peers = self.squeak_controller.get_peers()
diff --git a/squeaknode/admin/squeak_admin_server_servicer.py b/squeaknode/admin/squeak_admin_server_servicer.py
index 34368b8b..37406766 100644
--- a/squeaknode/admin/squeak_admin_server_servicer.py
+++ b/squeaknode/admin/squeak_admin_server_servicer.py
@@ -184,6 +184,9 @@ class SqueakAdminServerServicer(squeak_admin_pb2_grpc.SqueakAdminServicer):
             return squeak_admin_pb2.GetPeerReply()
         return reply
 
+    def GetPeerByAddress(self, request, context):
+        return self.handler.handle_get_squeak_peer_by_address(request)
+
     def GetPeers(self, request, context):
         return self.handler.handle_get_squeak_peers(request)
 
diff --git a/squeaknode/node/squeak_controller.py b/squeaknode/node/squeak_controller.py
index e5fede45..2da2e4a2 100644
--- a/squeaknode/node/squeak_controller.py
+++ b/squeaknode/node/squeak_controller.py
@@ -369,6 +369,9 @@ class SqueakController:
     def get_peer(self, peer_id: int) -> Optional[SqueakPeer]:
         return self.squeak_db.get_peer(peer_id)
 
+    def get_peer_by_address(self, peer_address: PeerAddress) -> Optional[SqueakPeer]:
+        return self.squeak_db.get_peer_by_address(peer_address)
+
     def get_peers(self):
         return self.squeak_db.get_peers()