Fixed #17671 - Cursors are now context managers.

1 files changed, 25 insertions, 0 deletions

diff --git a/tests/backends/tests.py b/tests/backends/tests.py
index a6badac9a9..73b17b0b81 100644
--- a/tests/backends/tests.py
+++ b/tests/backends/tests.py
@@ -613,6 +613,31 @@ class BackendTestCase(TestCase):
         with self.assertRaises(DatabaseError):
             cursor.execute(query)
 
+    def test_cursor_contextmanager(self):
+        """
+        Test that cursors can be used as a context manager
+        """
+        with connection.cursor() as cursor:
+            from django.db.backends.util import CursorWrapper
+            self.assertTrue(isinstance(cursor, CursorWrapper))
+        # Both InterfaceError and ProgrammingError seem to be used when
+        # accessing closed cursor (psycopg2 has InterfaceError, rest seem
+        # to use ProgrammingError).
+        with self.assertRaises(connection.features.closed_cursor_error_class):
+            # cursor should be closed, so no queries should be possible.
+            cursor.execute("select 1")
+
+    @unittest.skipUnless(connection.vendor == 'postgresql',
+                         "Psycopg2 specific cursor.closed attribute needed")
+    def test_cursor_contextmanager_closing(self):
+        # There isn't a generic way to test that cursors are closed, but
+        # psycopg2 offers us a way to check that by closed attribute.
+        # So, run only on psycopg2 for that reason.
+        with connection.cursor() as cursor:
+            from django.db.backends.util import CursorWrapper
+            self.assertTrue(isinstance(cursor, CursorWrapper))
+        self.assertTrue(cursor.closed)
+
 
 # We don't make these tests conditional because that means we would need to
 # check and differentiate between: