Get connection from yielding context manager of connection manager (#1351)

squeaknode/network/connection.py

@@ -21,7 +21,6 @@
 # SOFTWARE.
 import logging
 import threading
-from contextlib import contextmanager
 
 from squeak.messages import msg_addr
 from squeak.messages import msg_getaddr
@@ -69,19 +68,6 @@ class Connection(object):
             str(self.peer),
         )
 
-    @contextmanager
-    def connect(self, connection_manager):
-        logger.debug("Adding peer.")
-        connection_manager.add_peer(self.peer)
-        try:
-            logger.debug("Yielding peer.")
-            yield self
-        except Exception:
-            logger.exception("Peer connection failed.")
-        finally:
-            logger.debug("Removing peer.")
-            connection_manager.remove_peer(self.peer)
-
     def shutdown(self):
         logger.debug("Peet shutting down...")
         self.peer.stop()

squeaknode/network/connection_manager.py

@@ -21,11 +21,13 @@
 # SOFTWARE.
 import logging
 import threading
+from contextlib import contextmanager
 from typing import Dict
 from typing import List
 from typing import Optional
 
 from squeaknode.core.peer_address import PeerAddress
+from squeaknode.network.connection import Connection
 from squeaknode.network.peer import Peer
 from squeaknode.node.listener_subscription_client import EventListener
 
@@ -49,6 +51,19 @@ class ConnectionManager(object):
         self.single_peer_changed_listener = EventListener()
         self.accept_connections = True
 
+    @contextmanager
+    def connect(self, peer, squeak_controller):
+        logger.debug("Adding peer.")
+        self.add_peer(peer)
+        try:
+            logger.debug("Yielding connection.")
+            yield Connection(peer, squeak_controller)
+        except Exception:
+            logger.exception("Peer connection failed.")
+        finally:
+            logger.debug("Removing peer.")
+            self.remove_peer(peer)
+
     @property
     def peers(self) -> List[Peer]:
         return list(self._peers.values())

squeaknode/network/peer_handler.py

@@ -24,7 +24,6 @@ import socket
 import threading
 
 from squeaknode.core.peer_address import PeerAddress
-from squeaknode.network.connection import Connection
 from squeaknode.network.peer import Peer
 
 
@@ -100,9 +99,7 @@ class PeerHandler():
     def start_connection(self, peer: Peer):
         """Start a connection
         """
-        with Connection(peer, self.squeak_controller).connect(
-                self.connection_manager
-        ) as connection:
+        with self.connection_manager.connect(peer, self.squeak_controller) as connection:
             connection.handle_connection()