From 92725762f68bdf05309855e3642eb45a20efcce3 Mon Sep 17 00:00:00 2001
From: Jonathan Zernik <yzernik@gmail.com>
Date: Wed, 11 Aug 2021 20:28:00 -0700
Subject: [PATCH] Combine two context managers for peer connection (#914)

* Combine two context managers for peer connection

* Remove old comments
---
 squeaknode/network/peer.py         | 30 ++----------------------------
 squeaknode/network/peer_handler.py | 13 ++++---------
 2 files changed, 6 insertions(+), 37 deletions(-)

diff --git a/squeaknode/network/peer.py b/squeaknode/network/peer.py
index e6519b47..b2d08faa 100644
--- a/squeaknode/network/peer.py
+++ b/squeaknode/network/peer.py
@@ -160,10 +160,6 @@ class Peer(object):
         logger.info('Received msg {} from {}'.format(msg, self))
         return msg
 
-    # def stop(self):
-    #     logger.info("Stopping peer: {}".format(self))
-    #     self.stopped.set()
-
     def close(self):
         logger.info("closing peer socket: {}".format(self._peer_socket))
         try:
@@ -183,20 +179,6 @@ class Peer(object):
             logger.info('Failed to send msg to {}'.format(self))
             self.close()
 
-    # def __enter__(self):
-    #     logger.debug('Setting up peer {} ...'.format(self))
-    #     msg_receiver = MessageReceiver(
-    #         self._peer_socket, self._recv_msg_queue, self.stopped)
-    #     threading.Thread(
-    #         target=msg_receiver.recv_msgs,
-    #         args=(),
-    #     ).start()
-    #     return self
-
-    # def __exit__(self, *exc):
-    #     self.close()
-    #     logger.debug('Closed connection to peer {} ...'.format(self))
-
     def handshake(self, squeak_controller):
         if self.outgoing:
             local_version = self.version_pkt(squeak_controller)
@@ -244,7 +226,7 @@ class Peer(object):
         logger.info('Finished handling messages...')
 
     @contextmanager
-    def start_peer(self):
+    def open_connection(self, squeak_controller):
         logger.debug('Setting up peer {} ...'.format(self))
         try:
             msg_receiver = MessageReceiver(
@@ -253,20 +235,12 @@ class Peer(object):
                 target=msg_receiver.recv_msgs,
                 args=(),
             ).start()
+            self.handshake(squeak_controller)
             yield self
         finally:
             self.close()
             logger.debug('Closed connection to peer {} ...'.format(self))
 
-    @contextmanager
-    def open_connection(self, squeak_controller):
-        logger.info(
-            'Starting handshake connection with peer ... {}'.format(self))
-        self.handshake(squeak_controller)
-        yield self
-        logger.info(
-            'Finished handshake connection with peer ... {}'.format(self))
-
     def __repr__(self):
         return "Peer(%s)" % (self.address_string)
 
diff --git a/squeaknode/network/peer_handler.py b/squeaknode/network/peer_handler.py
index f8a4fce0..a681c9b3 100644
--- a/squeaknode/network/peer_handler.py
+++ b/squeaknode/network/peer_handler.py
@@ -32,16 +32,11 @@ class PeerHandler():
 
         logger.debug(
             'Setting up controller for peer address {} ...'.format(address))
-        logger.info(
-            'Setting up controller for peer address {} ...'.format(address))
-        with Peer(peer_socket, address, outgoing).start_peer() as p:
-            with p.open_connection(self.squeak_controller) as p2:
-                # p.stopped.wait()
-                self.connection_manager.add_peer(p2)
-                p2.handle_messages(self.squeak_controller)
-        self.connection_manager.remove_peer(p2)
+        with Peer(peer_socket, address, outgoing).open_connection(self.squeak_controller) as peer:
+            self.connection_manager.add_peer(peer)
+            peer.handle_messages(self.squeak_controller)
+        self.connection_manager.remove_peer(peer)
         logger.debug('Stopped controller for peer address {}.'.format(address))
-        logger.info('Stopped controller for peer address {}.'.format(address))
 
     def handle_connection(self, peer_socket, address, outgoing):
         threading.Thread(