Ensure no acquire after close for ConnectionPool

Author: technige

neo4j/v1/bolt.py

@@ -378,6 +378,8 @@ class ConnectionPool(object):
     """ A collection of connections to one or more server addresses.
     """
 
+    closed = False
+
     def __init__(self, connector):
         self.connector = connector
         self.connections = {}
@@ -393,6 +395,9 @@ def acquire(self, address):
         """ Acquire a connection to a given address from the pool.
         This method is thread safe.
         """
+        if self.closed:
+            raise ServiceUnavailable("This connection pool is closed so no new "
+                                     "connections may be acquired")
         with self.lock:
             try:
                 connections = self.connections[address]
@@ -433,6 +438,7 @@ def close(self):
         This method is thread safe.
         """
         with self.lock:
+            self.closed = True
             for address in list(self.connections):
                 self.remove(address)
 

test/test_connection.py

@@ -20,7 +20,7 @@
 
 from socket import create_connection
 
-from neo4j.v1 import ConnectionPool
+from neo4j.v1 import ConnectionPool, ServiceUnavailable
 
 from test.util import ServerTestCase
 
@@ -98,3 +98,9 @@ def test_releasing_twice(self):
         self.assert_pool_size(address, 0, 1)
         self.pool.release(connection)
         self.assert_pool_size(address, 0, 1)
+
+    def test_cannot_acquire_after_close(self):
+        with ConnectionPool(lambda a: QuickConnection(create_connection(a))) as pool:
+            pool.close()
+            with self.assertRaises(ServiceUnavailable):
+                _ = pool.acquire("X")