Fix tests

Author: abhinavsingh

proxy/core/connection/pool.py

@@ -77,38 +77,32 @@ def __init__(self) -> None:
         self.connections: Dict[int, TcpServerConnection] = {}
         self.pools: Dict[Tuple[str, int], Set[TcpServerConnection]] = {}
 
-    def add(self, addr: Tuple[str, int]) -> TcpServerConnection:
-        """Creates and add a new connection to the pool."""
-        new_conn = TcpServerConnection(addr[0], addr[1])
-        new_conn.connect()
-        self._add(new_conn)
-        logger.debug(
-            'Created new connection#{2} for upstream {0}:{1}'.format(
-                addr[0], addr[1], id(new_conn),
-            ),
-        )
-        return new_conn
-
     def acquire(self, addr: Tuple[str, int]) -> Tuple[bool, TcpServerConnection]:
         """Returns a reusable connection from the pool.
 
         If none exists, will create and return a new connection."""
+        created, conn = False, None
         if addr in self.pools:
             for old_conn in self.pools[addr]:
                 if old_conn.is_reusable():
-                    old_conn.mark_inuse()
+                    conn = old_conn
                     logger.debug(
                         'Reusing connection#{2} for upstream {0}:{1}'.format(
                             addr[0], addr[1], id(old_conn),
                         ),
                     )
-                    return False, old_conn
-        new_conn = self.add(addr)
-        new_conn.mark_inuse()
-        return True, new_conn
+                    break
+        if conn is None:
+            created, conn = True, self.add(addr)
+        conn.mark_inuse()
+        return created, conn
 
     def release(self, conn: TcpServerConnection) -> None:
-        """Release a previously acquired connection."""
+        """Release a previously acquired connection.
+
+        Releasing a connection will shutdown and close the socket
+        including internal pool cleanup.
+        """
         assert not conn.is_reusable()
         logger.debug(
             'Removing connection#{2} from pool from upstream {0}:{1}'.format(
@@ -118,7 +112,8 @@ def release(self, conn: TcpServerConnection) -> None:
         self._remove(conn.connection.fileno())
 
     def retain(self, conn: TcpServerConnection) -> None:
-        """Retained the connection in the pool for reusability."""
+        """Retained previously acquired connection in the pool for reusability."""
+        assert not conn.closed
         logger.debug(
             'Retaining connection#{2} to upstream {0}:{1}'.format(
                 conn.addr[0], conn.addr[1], id(conn),
@@ -150,6 +145,23 @@ async def handle_events(self, readables: Readables, _writables: Writables) -> bo
             self._remove(fileno)
         return False
 
+    def add(self, addr: Tuple[str, int]) -> TcpServerConnection:
+        """Creates, connects and adds a new connection to the pool.
+
+        Returns newly created connection.
+
+        NOTE: You must not use the returned connection, instead use `acquire`.
+        """
+        new_conn = TcpServerConnection(addr[0], addr[1])
+        new_conn.connect()
+        self._add(new_conn)
+        logger.debug(
+            'Created new connection#{2} for upstream {0}:{1}'.format(
+                addr[0], addr[1], id(new_conn),
+            ),
+        )
+        return new_conn
+
     def _add(self, conn: TcpServerConnection) -> None:
         """Adds a new connection to internal data structure."""
         if conn.addr not in self.pools:

tests/core/test_conn_pool.py

@@ -27,24 +27,28 @@ def test_acquire_and_retain_and_reacquire(self, mock_tcp_server_connection: mock
         mock_conn = mock_tcp_server_connection.return_value
         addr = mock_conn.addr
         mock_conn.is_reusable.side_effect = [
-            False, True, True,
+            True,
         ]
         mock_conn.closed = False
         # Acquire
         created, conn = pool.acquire(addr)
-        self.assertTrue(created)
         mock_tcp_server_connection.assert_called_once_with(addr[0], addr[1])
+        mock_conn.mark_inuse.assert_called_once()
+        mock_conn.reset.assert_not_called()
+        self.assertTrue(created)
         self.assertEqual(conn, mock_conn)
         self.assertEqual(len(pool.pools[addr]), 1)
         self.assertTrue(conn in pool.pools[addr])
-        # Release (connection must be retained because not closed)
+        self.assertEqual(len(pool.connections), 1)
+        self.assertEqual(pool.connections[conn.connection.fileno()], mock_conn)
+        # Retail
         pool.retain(conn)
         self.assertEqual(len(pool.pools[addr]), 1)
         self.assertTrue(conn in pool.pools[addr])
+        mock_conn.reset.assert_called_once()
         # Reacquire
         created, conn = pool.acquire(addr)
         self.assertFalse(created)
-        mock_conn.reset.assert_called_once()
         self.assertEqual(conn, mock_conn)
         self.assertEqual(len(pool.pools[addr]), 1)
         self.assertTrue(conn in pool.pools[addr])