Close DB connections when testing async functionality.

Author: tim-schilling

tests/panels/test_sql.py

@@ -18,6 +18,14 @@
 
 from ..base import BaseTestCase
 from ..models import PostgresJSON
+from ..sync import database_sync_to_async
+
+
+def sql_call(use_iterator=False):
+    qs = User.objects.all()
+    if use_iterator:
+        qs = qs.iterator()
+    return list(qs)
 
 
 class SQLPanelTestCase(BaseTestCase):
@@ -32,7 +40,7 @@ def test_disabled(self):
     def test_recording(self):
         self.assertEqual(len(self.panel._queries), 0)
 
-        list(User.objects.all())
+        sql_call()
 
         # ensure query was logged
         self.assertEqual(len(self.panel._queries), 1)
@@ -51,7 +59,7 @@ def test_recording(self):
     def test_recording_chunked_cursor(self):
         self.assertEqual(len(self.panel._queries), 0)
 
-        list(User.objects.all().iterator())
+        sql_call(use_iterator=True)
 
         # ensure query was logged
         self.assertEqual(len(self.panel._queries), 1)
@@ -61,7 +69,7 @@ def test_recording_chunked_cursor(self):
         wraps=sql_tracking.NormalCursorWrapper,
     )
     def test_cursor_wrapper_singleton(self, mock_wrapper):
-        list(User.objects.all())
+        sql_call()
 
         # ensure that cursor wrapping is applied only once
         self.assertEqual(mock_wrapper.call_count, 1)
@@ -71,7 +79,7 @@ def test_cursor_wrapper_singleton(self, mock_wrapper):
         wraps=sql_tracking.NormalCursorWrapper,
     )
     def test_chunked_cursor_wrapper_singleton(self, mock_wrapper):
-        list(User.objects.all().iterator())
+        sql_call(use_iterator=True)
 
         # ensure that cursor wrapping is applied only once
         self.assertEqual(mock_wrapper.call_count, 1)
@@ -81,7 +89,7 @@ def test_chunked_cursor_wrapper_singleton(self, mock_wrapper):
         wraps=sql_tracking.NormalCursorWrapper,
     )
     async def test_cursor_wrapper_async(self, mock_wrapper):
-        await sync_to_async(list)(User.objects.all())
+        await sync_to_async(sql_call)()
 
         self.assertEqual(mock_wrapper.call_count, 1)
 
@@ -91,11 +99,13 @@ async def test_cursor_wrapper_async(self, mock_wrapper):
     )
     async def test_cursor_wrapper_asyncio_ctx(self, mock_wrapper):
         self.assertTrue(sql_tracking.recording.get())
-        await sync_to_async(list)(User.objects.all())
+        await sync_to_async(sql_call)()
 
         async def task():
             sql_tracking.recording.set(False)
-            await sync_to_async(list)(User.objects.all())
+            # Calling this in another context requires the db connections
+            # to be closed properly.
+            await database_sync_to_async(sql_call)()
 
         # Ensure this is called in another context
         await asyncio.create_task(task())
@@ -106,7 +116,7 @@ async def task():
     def test_generate_server_timing(self):
         self.assertEqual(len(self.panel._queries), 0)
 
-        list(User.objects.all())
+        sql_call()
 
         response = self.panel.process_request(self.request)
         self.panel.generate_stats(self.request, response)
@@ -372,7 +382,7 @@ def test_disable_stacktraces(self):
         self.assertEqual(len(self.panel._queries), 0)
 
         with self.settings(DEBUG_TOOLBAR_CONFIG={"ENABLE_STACKTRACES": False}):
-            list(User.objects.all())
+            sql_call()
 
         # ensure query was logged
         self.assertEqual(len(self.panel._queries), 1)

tests/sync.py

@@ -0,0 +1,22 @@
+"""
+Taken from channels.db
+"""
+from asgiref.sync import SyncToAsync
+from django.db import close_old_connections
+
+
+class DatabaseSyncToAsync(SyncToAsync):
+    """
+    SyncToAsync version that cleans up old database connections when it exits.
+    """
+
+    def thread_handler(self, loop, *args, **kwargs):
+        close_old_connections()
+        try:
+            return super().thread_handler(loop, *args, **kwargs)
+        finally:
+            close_old_connections()
+
+
+# The class is TitleCased, but we want to encourage use as a callable/decorator
+database_sync_to_async = DatabaseSyncToAsync