From 1af6607fc38c6b20aaae1a45bd3f86b6088cae64 Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki.linnakangas@iki.fi>
Date: Thu, 27 May 2021 12:59:45 +0300
Subject: [PATCH] Add a test for restarting and recreating compute node.

This is working; let's keep it that way. This also adds test coverage for
the 'zenith pg stop --destroy' option that was added in commit 6ad6e5bd.
---
 .../batch_others/test_restart_compute.py      | 59 +++++++++++++++++++
 test_runner/fixtures/zenith_fixtures.py       | 19 +++++-
 2 files changed, 75 insertions(+), 3 deletions(-)
 create mode 100644 test_runner/batch_others/test_restart_compute.py

diff --git a/test_runner/batch_others/test_restart_compute.py b/test_runner/batch_others/test_restart_compute.py
new file mode 100644
index 0000000000..b6e13b3f36
--- /dev/null
+++ b/test_runner/batch_others/test_restart_compute.py
@@ -0,0 +1,59 @@
+import pytest
+import getpass
+import psycopg2
+import time
+
+pytest_plugins = ("fixtures.zenith_fixtures")
+
+#
+# Test restarting and recreating a postgres instance
+#
+def test_restart_compute(zenith_cli, pageserver, postgres, pg_bin):
+    zenith_cli.run(["branch", "test_restart_compute", "empty"]);
+
+    pg = postgres.create_start('test_restart_compute')
+    print("postgres is running on 'test_restart_compute' branch")
+
+    pg_conn = psycopg2.connect(pg.connstr());
+    pg_conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+    cur = pg_conn.cursor()
+
+    # Create table, and insert a row
+    cur.execute('CREATE TABLE foo (t text)');
+    cur.execute("INSERT INTO foo VALUES ('bar')");
+
+    # Stop and restart the Postgres instance
+    pg_conn.close();
+    pg.stop();
+    pg.start();
+    pg_conn = psycopg2.connect(pg.connstr());
+    pg_conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+    cur = pg_conn.cursor()
+
+    # We can still see the row
+    cur.execute('SELECT count(*) FROM foo');
+    assert(cur.fetchone()[0] == 1);
+
+    # Insert another row
+    cur.execute("INSERT INTO foo VALUES ('bar2')");
+    cur.execute('SELECT count(*) FROM foo');
+    assert(cur.fetchone()[0] == 2);
+
+    # FIXME: Currently, there is no guarantee that by the time the INSERT commits, the WAL
+    # has been streamed safely to the WAL safekeeper or page server. It is merely stored
+    # on the Postgres instance's local disk. Sleep a little, to give it time to be
+    # streamed. This should be removed, when we have the ability to run the Postgres
+    # instance -> safekeeper streaming in synchronous mode.
+    time.sleep(5)
+
+    # Stop, and destroy the Postgres instance. Then recreate and restart it.
+    pg_conn.close();
+    pg.stop_and_destroy();
+    pg.create_start('test_restart_compute');
+    pg_conn = psycopg2.connect(pg.connstr());
+    pg_conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+    cur = pg_conn.cursor()
+
+    # We can still see the rows
+    cur.execute('SELECT count(*) FROM foo');
+    assert(cur.fetchone()[0] == 2);
diff --git a/test_runner/fixtures/zenith_fixtures.py b/test_runner/fixtures/zenith_fixtures.py
index 389120d32e..1980d3ca1c 100644
--- a/test_runner/fixtures/zenith_fixtures.py
+++ b/test_runner/fixtures/zenith_fixtures.py
@@ -168,14 +168,18 @@ class Postgres:
         self.branch = None
         # path to conf is <repo_dir>/pgdatadirs/<branch_name>/postgresql.conf
 
-    def create_start(self, branch, config_lines=None):
-        """ create the pg data directory, and start the server """
+    def create(self, branch, config_lines=None):
+        """ create the pg data directory """
         self.zenith_cli.run(['pg', 'create', branch])
         self.branch = branch
         if config_lines is None:
             config_lines = []
         self.config(config_lines)
-        self.zenith_cli.run(['pg', 'start', branch])
+        return
+
+    def start(self):
+        """ start the server """
+        self.zenith_cli.run(['pg', 'start', self.branch])
         self.running = True
         return
 
@@ -189,9 +193,18 @@ class Postgres:
                 conf.write('\n')
 
     def stop(self):
+        """ stop the server """
         if self.running:
             self.zenith_cli.run(['pg', 'stop', self.branch])
 
+    def stop_and_destroy(self):
+        self.zenith_cli.run(['pg', 'stop', '--destroy', self.branch])
+
+    def create_start(self, branch, config_lines=None):
+        self.create(branch, config_lines);
+        self.start();
+        return
+
     # Return a libpq connection string to connect to the Postgres instance
     def connstr(self, dbname='postgres'):
         conn_str = 'host={} port={} dbname={} user={}'.format(