Comments, more unittests

Author: 1st1

asyncpg/_testbase.py

@@ -190,16 +190,32 @@ def start_cluster(cls, ClusterCls, *,
         return _start_cluster(ClusterCls, cluster_kwargs, server_settings)
 
 
+def with_connection_options(**options):
+    if not options:
+        raise ValueError('no connection options were specified')
+
+    def wrap(func):
+        func.__connect_options__ = options
+        return func
+
+    return wrap
+
+
 class ConnectedTestCase(ClusterTestCase):
 
     def getExtraConnectOptions(self):
         return {}
 
     def setUp(self):
         super().setUp()
-        opts = self.getExtraConnectOptions()
+
+        # Extract options set up with `with_connection_options`.
+        test_func = getattr(self, self._testMethodName).__func__
+        opts = getattr(test_func, '__connect_options__', {})
+
         self.con = self.loop.run_until_complete(
             self.cluster.connect(database='postgres', loop=self.loop, **opts))
+
         self.server_version = self.con.get_server_version()
 
     def tearDown(self):

asyncpg/connection.py

@@ -880,10 +880,13 @@ def set_max_lifetime(self, new_lifetime):
         assert new_lifetime >= 0
         self._max_lifetime = new_lifetime
         for entry in self._entries.values():
+            # For every entry cancel the existing callback
+            # and setup a new one if necessary.
             self._set_entry_timeout(entry)
 
     def get(self, query, *, promote=True):
         if not self._max_size:
+            # The cache is disabled.
             return
 
         entry = self._entries.get(query)  # type: _StatementCacheEntry
@@ -898,6 +901,7 @@ def get(self, query, *, promote=True):
             return
 
         if promote:
+            # `promote` is `False` when `get()` is called by `has()`.
             self._entries.move_to_end(query, last=True)
 
         return entry._statement
@@ -907,21 +911,31 @@ def has(self, query):
 
     def put(self, query, statement):
         if not self._max_size:
+            # The cache is disabled.
             return
 
         self._entries[query] = self._new_entry(query, statement)
+
+        # Check if the cache is bigger than max_size and trim it
+        # if necessary.
         self._maybe_cleanup()
 
     def iter_statements(self):
         return (e._statement for e in self._entries.values())
 
     def clear(self):
+        # First, make sure that we cancel all scheduled callbacks.
         for entry in self._entries.values():
             self._clear_entry_callback(entry)
+
+        # Clear the entries dict.
         self._entries.clear()
 
     def _set_entry_timeout(self, entry):
+        # Clear the existing timeout.
         self._clear_entry_callback(entry)
+
+        # Set the new timeout if it's not 0.
         if self._max_lifetime:
             entry._cleanup_cb = self._loop.call_later(
                 self._max_lifetime, self._on_entry_expired, entry)
@@ -943,9 +957,13 @@ def _clear_entry_callback(self, entry):
             entry._cleanup_cb.cancel()
 
     def _maybe_cleanup(self):
+        # Delete cache entries until the size of the cache is `max_size`.
         while len(self._entries) > self._max_size:
             old_query, old_entry = self._entries.popitem(last=False)
             self._clear_entry_callback(old_entry)
+
+            # Let the connection know that the statement was removed
+            # from the cache.
             self._on_remove(old_entry._statement)
 
 

tests/test_prepare.py

@@ -430,10 +430,9 @@ async def test_prepare_statement_invalid(self):
         finally:
             await self.con.execute('DROP TABLE tab1')
 
+    @tb.with_connection_options(statement_cache_size=0)
     async def test_prepare_23_no_stmt_cache_seq(self):
-        # Disable cache, which will force connections to use
-        # anonymous prepared statements.
-        self.con._stmt_cache.set_max_size(0)
+        self.assertEqual(self.con._stmt_cache.get_max_size(), 0)
 
         async def check_simple():
             # Run a simple query a few times.
@@ -462,8 +461,11 @@ async def check_simple():
         # operation.
         await check_simple()
 
+    @tb.with_connection_options(max_cached_statement_lifetime=142)
     async def test_prepare_24_max_lifetime(self):
         cache = self.con._stmt_cache
+
+        self.assertEqual(cache.get_max_lifetime(), 142)
         cache.set_max_lifetime(1)
 
         s = await self.con.prepare('SELECT 1')
@@ -479,3 +481,35 @@ async def test_prepare_24_max_lifetime(self):
 
         s = await self.con.prepare('SELECT 1')
         self.assertIsNot(s._state, state)
+
+    @tb.with_connection_options(max_cached_statement_lifetime=0.5)
+    async def test_prepare_25_max_lifetime_reset(self):
+        cache = self.con._stmt_cache
+
+        s = await self.con.prepare('SELECT 1')
+        state = s._state
+
+        # Disable max_lifetime
+        cache.set_max_lifetime(0)
+
+        await asyncio.sleep(1, loop=self.loop)
+
+        # The statement should still be cached (as we disabled the timeout).
+        s = await self.con.prepare('SELECT 1')
+        self.assertIs(s._state, state)
+
+    @tb.with_connection_options(max_cached_statement_lifetime=0.5)
+    async def test_prepare_26_max_lifetime_max_size(self):
+        cache = self.con._stmt_cache
+
+        s = await self.con.prepare('SELECT 1')
+        state = s._state
+
+        # Disable max_lifetime
+        cache.set_max_size(0)
+
+        s = await self.con.prepare('SELECT 1')
+        self.assertIsNot(s._state, state)
+
+        # Check that nothing crashes after the initial timeout
+        await asyncio.sleep(1, loop=self.loop)

tests/test_timeout.py

@@ -128,11 +128,7 @@ async def test_invalid_timeout(self):
 
 class TestConnectionCommandTimeout(tb.ConnectedTestCase):
 
-    def getExtraConnectOptions(self):
-        return {
-            'command_timeout': 0.2
-        }
-
+    @tb.with_connection_options(command_timeout=0.2)
     async def test_command_timeout_01(self):
         for methname in {'fetch', 'fetchrow', 'fetchval', 'execute'}:
             with self.assertRaises(asyncio.TimeoutError), \
@@ -151,12 +147,8 @@ async def _get_statement(self, query, timeout):
 
 class TestTimeoutCoversPrepare(tb.ConnectedTestCase):
 
-    def getExtraConnectOptions(self):
-        return {
-            '__connection_class__': SlowPrepareConnection,
-            'command_timeout': 0.3
-        }
-
+    @tb.with_connection_options(__connection_class__=SlowPrepareConnection,
+                                command_timeout=0.3)
     async def test_timeout_covers_prepare_01(self):
         for methname in {'fetch', 'fetchrow', 'fetchval', 'execute'}:
             with self.assertRaises(asyncio.TimeoutError):