From 72f6221f4359a74b3bbba619c4cbd37d6912f7dd Mon Sep 17 00:00:00 2001
From: ankiaga <ankiaga@google.com>
Date: Wed, 6 Dec 2023 16:49:21 +0530
Subject: [PATCH] Added tests for exception in commit and rollback

---
 tests/system/test_dbapi.py | 45 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/tests/system/test_dbapi.py b/tests/system/test_dbapi.py
index 60aa772c24..89b1068793 100644
--- a/tests/system/test_dbapi.py
+++ b/tests/system/test_dbapi.py
@@ -126,6 +126,51 @@ def test_commit(self, client_side):
 
         assert got_rows == [updated_row]
 
+    @pytest.mark.parametrize("client_side", [True, False])
+    def test_commit_exception(self, client_side):
+        """Test that if exception during commit method is caught, then
+        subsequent operations on same Cursor and Connection object works
+        properly."""
+        self._execute_common_statements(self._cursor)
+        # deleting the session to fail the commit
+        self._conn._session.delete()
+        try:
+            if client_side:
+                self._cursor.execute("""COMMIT""")
+            else:
+                self._conn.commit()
+        except Exception:
+            pass
+
+        # Testing that the connection and Cursor are in proper state post commit
+        updated_row = self._execute_common_statements(self._cursor)
+        self._cursor.execute("SELECT * FROM contacts")
+        got_rows = self._cursor.fetchall()
+        self._conn.commit()
+
+        assert got_rows == [updated_row]
+
+    def test_rollback_exception(self, client_side):
+        """Test that if exception during rollback method is caught, then
+        subsequent operations on same Cursor and Connection object works
+        properly."""
+        self._execute_common_statements(self._cursor)
+        # deleting the session to fail the rollback
+        self._conn._session.delete()
+        try:
+            self._conn.rollback()
+        except Exception:
+            pass
+
+        # Testing that the connection and Cursor are in proper state post
+        # exception in rollback
+        updated_row = self._execute_common_statements(self._cursor)
+        self._cursor.execute("SELECT * FROM contacts")
+        got_rows = self._cursor.fetchall()
+        self._conn.commit()
+
+        assert got_rows == [updated_row]
+
     def test_cursor_execute_exception(self):
         """Test that if exception in Cursor's execute method is caught when
         Connection is not in autocommit mode, then subsequent operations on