From 2679194c56e7cd43cc36406aa6fe2e4783bc536b Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Thu, 9 Dec 2021 14:27:33 +0100
Subject: [PATCH] Add tests to verify out-of-order detections without threads

---
 tests/test_transaction.py       | 45 +++++++++++++++++++++++++++++++++
 tests/test_transaction_async.py | 36 +++++++++++++++++++++++++-
 2 files changed, 80 insertions(+), 1 deletion(-)

diff --git a/tests/test_transaction.py b/tests/test_transaction.py
index 07a05bcbd..bc7c374f9 100644
--- a/tests/test_transaction.py
+++ b/tests/test_transaction.py
@@ -1,3 +1,4 @@
+import sys
 import logging
 from threading import Thread, Event
 
@@ -55,6 +56,16 @@ def in_transaction(conn):
         assert False, conn.pgconn.transaction_status
 
 
+def get_exc_info(exc):
+    """Return the exc info for an exception or a success if exc is None"""
+    if not exc:
+        return (None,) * 3
+    try:
+        raise exc
+    except exc:
+        return sys.exc_info()
+
+
 class ExpectedException(Exception):
     pass
 
@@ -657,6 +668,40 @@ def test_str(conn):
     assert "(terminated)" in str(tx)
 
 
+@pytest.mark.parametrize("exit_error", [None, ZeroDivisionError, Rollback])
+def test_out_of_order_exit(conn, exit_error):
+    conn.autocommit = True
+
+    t1 = conn.transaction()
+    t1.__enter__()
+
+    t2 = conn.transaction()
+    t2.__enter__()
+
+    with pytest.raises(ProgrammingError):
+        t1.__exit__(*get_exc_info(exit_error))
+
+    with pytest.raises(ProgrammingError):
+        t2.__exit__(*get_exc_info(exit_error))
+
+
+@pytest.mark.parametrize("exit_error", [None, ZeroDivisionError, Rollback])
+def test_out_of_order_implicit_begin(conn, exit_error):
+    conn.execute("select 1")
+
+    t1 = conn.transaction()
+    t1.__enter__()
+
+    t2 = conn.transaction()
+    t2.__enter__()
+
+    with pytest.raises(ProgrammingError):
+        t1.__exit__(*get_exc_info(exit_error))
+
+    with pytest.raises(ProgrammingError):
+        t2.__exit__(*get_exc_info(exit_error))
+
+
 @pytest.mark.parametrize("what", ["commit", "rollback", "error"])
 def test_concurrency(conn, what):
     conn.autocommit = True
diff --git a/tests/test_transaction_async.py b/tests/test_transaction_async.py
index 4bfce1595..82df069bb 100644
--- a/tests/test_transaction_async.py
+++ b/tests/test_transaction_async.py
@@ -7,7 +7,7 @@ from psycopg import AsyncConnection, ProgrammingError, Rollback
 from psycopg._compat import create_task
 
 from .test_transaction import in_transaction, insert_row, inserted
-from .test_transaction import ExpectedException
+from .test_transaction import ExpectedException, get_exc_info
 from .test_transaction import create_test_table  # noqa  # autouse fixture
 
 pytestmark = pytest.mark.asyncio
@@ -625,6 +625,40 @@ async def test_str(aconn):
     assert "(terminated)" in str(tx)
 
 
+@pytest.mark.parametrize("exit_error", [None, ZeroDivisionError, Rollback])
+async def test_out_of_order_exit(aconn, exit_error):
+    await aconn.set_autocommit(True)
+
+    t1 = aconn.transaction()
+    await t1.__aenter__()
+
+    t2 = aconn.transaction()
+    await t2.__aenter__()
+
+    with pytest.raises(ProgrammingError):
+        await t1.__aexit__(*get_exc_info(exit_error))
+
+    with pytest.raises(ProgrammingError):
+        await t2.__aexit__(*get_exc_info(exit_error))
+
+
+@pytest.mark.parametrize("exit_error", [None, ZeroDivisionError, Rollback])
+async def test_out_of_order_implicit_begin(aconn, exit_error):
+    await aconn.execute("select 1")
+
+    t1 = aconn.transaction()
+    await t1.__aenter__()
+
+    t2 = aconn.transaction()
+    await t2.__aenter__()
+
+    with pytest.raises(ProgrammingError):
+        await t1.__aexit__(*get_exc_info(exit_error))
+
+    with pytest.raises(ProgrammingError):
+        await t2.__aexit__(*get_exc_info(exit_error))
+
+
 @pytest.mark.parametrize("what", ["commit", "rollback", "error"])
 async def test_concurrency(aconn, what):
     await aconn.set_autocommit(True)
-- 
2.47.2