From e8f0a9bb801ed7b2ee7c16e2f3edb0e32700aa3d Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki.linnakangas@iki.fi>
Date: Tue, 25 May 2021 11:11:32 +0300
Subject: [PATCH] Add test for prepared transactions.

---
 test_runner/batch_others/test_twophase.py | 50 +++++++++++++++++++++++
 1 file changed, 50 insertions(+)
 create mode 100644 test_runner/batch_others/test_twophase.py

diff --git a/test_runner/batch_others/test_twophase.py b/test_runner/batch_others/test_twophase.py
new file mode 100644
index 0000000000..679fc43b71
--- /dev/null
+++ b/test_runner/batch_others/test_twophase.py
@@ -0,0 +1,50 @@
+#
+# Test branching, when a transaction is in prepared state
+#
+import pytest
+import getpass
+import psycopg2
+
+pytest_plugins = ("fixtures.zenith_fixtures")
+
+def test_twophase(zenith_cli, pageserver, postgres, pg_bin):
+    zenith_cli.run(["branch", "test_twophase", "empty"]);
+
+    pg = postgres.create_start('test_twophase', ['max_prepared_transactions=5'])
+    print("postgres is running on 'test_twophase' branch")
+
+    conn = psycopg2.connect(pg.connstr());
+    conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+    cur = conn.cursor()
+
+    cur.execute('CREATE TABLE foo (t text)');
+
+    # Prepare a transaction that will insert a row
+    cur.execute('BEGIN');
+    cur.execute("INSERT INTO foo VALUES ('one')");
+    cur.execute("PREPARE TRANSACTION 'insert_one'");
+
+    # Prepare another transaction that will insert a row
+    cur.execute('BEGIN');
+    cur.execute("INSERT INTO foo VALUES ('two')");
+    cur.execute("PREPARE TRANSACTION 'insert_two'");
+
+    # Create a branch with the transaction in prepared state
+    zenith_cli.run(["branch", "test_twophase_prepared", "test_twophase"]);
+
+    pg2 = postgres.create_start('test_twophase_prepared', ['max_prepared_transactions=5'])
+    conn2 = psycopg2.connect(pg2.connstr());
+    conn2.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+    cur2 = conn2.cursor()
+
+    # On the new branch, commit one of the prepared transactions, abort the other one.
+    cur2.execute("COMMIT PREPARED 'insert_one'");
+    cur2.execute("ROLLBACK PREPARED 'insert_two'");
+
+    cur2.execute('SELECT * FROM foo');
+    assert(cur2.fetchall() == [('one',)]);
+
+    # Neither insert is visible on the original branch, the transactions are still
+    # in prepared state there.
+    cur.execute('SELECT * FROM foo');
+    assert(cur.fetchall() == []);