Merge branch 'master' of github.com:postgrespro/testgres

Author: funbringer

README.md

@@ -85,12 +85,23 @@ node.stop()
 
 It is essential to clean everything up, so make sure to call `node.cleanup()` once you've finished all of your tests.
 
+Nodes support python logging system, so if you have configured logging
+in your tests you can use it to redirect postgres logs to yours.
+To do that just use `use_logging` argument like here:
+
+```python
+node = testgres.get_new_node('master', use_logging=True)
+```
+
+Also you can find working configuration sample for logging in tests.
+
 Please see `testgres/tests` directory for replication configuration example.
 > Note: you could take a look at [`pg_pathman`](https://github.com/postgrespro/pg_pathman) to get an idea of `testgres`' capabilities.
 
 
 ## Authors
 
-Ildar Musin <i.musin(at)postgrespro.ru> Postgres Professional Ltd., Russia     
-Dmitry Ivanov <d.ivanov(at)postgrespro.ru> Postgres Professional Ltd., Russia   
-Ildus Kurbangaliev <i.kurbangaliev(at)postgrespro.ru> Postgres Professional Ltd., Russia
+[Ildar Musin](https://github.com/zilder) <i.musin(at)postgrespro.ru> Postgres Professional Ltd., Russia     
+[Dmitry Ivanov](https://github.com/funbringer) <d.ivanov(at)postgrespro.ru> Postgres Professional Ltd., Russia   
+[Ildus Kurbangaliev](https://github.com/ildus) <i.kurbangaliev(at)postgrespro.ru> Postgres Professional Ltd., Russia     
+[Yury Zhuravlev](https://github.com/stalkerg) <stalkerg(at)gmail.com>

setup.py

@@ -2,7 +2,7 @@
 setup(
     name='testgres',
     packages=['testgres'],
-    version='0.3.0',
+    version='0.3.1',
     description='Testing utility for postgresql and its extensions',
     license='PostgreSQL',
     author='Ildar Musin',

testgres/testgres.py

@@ -50,8 +50,10 @@
 
 registered_nodes = []
 util_threads = []
+tmp_dirs = []
 last_assigned_port = int(random.random() * 16384) + 49152
 pg_config_data = {}
+base_data_dir = None
 
 
 class ClusterException(Exception):
@@ -208,6 +210,7 @@ def __init__(self, name, port, base_dir=None, use_logging=False):
         self.port = port
         if base_dir is None:
             self.base_dir = tempfile.mkdtemp()
+            tmp_dirs.append(self.base_dir)
         else:
             self.base_dir = base_dir
         if not os.path.exists(self.logs_dir):
@@ -246,6 +249,31 @@ def get_bin_path(self, filename):
         else:
             return os.path.join(pg_config.get("BINDIR"), filename)
 
+    def initdb(self, directory, initdb_params=[]):
+        initdb = self.get_bin_path("initdb")
+        initdb_logfile = os.path.join(self.logs_dir, "initdb.log")
+
+        with open(initdb_logfile, 'a') as file_out:
+            ret = subprocess.call(
+                [initdb, directory, "-N"] + initdb_params,
+                stdout=file_out,
+                stderr=subprocess.STDOUT)
+
+            if ret:
+                raise ClusterException("Cluster initialization failed. You"
+                    " can find additional information at '%s'" % initdb_logfile)
+
+    def _setup_data_dir(self, data_dir):
+        global base_data_dir
+
+        if base_data_dir is None:
+            base_data_dir = tempfile.mkdtemp()
+            tmp_dirs.append(base_data_dir)
+            self.initdb(base_data_dir)
+
+        shutil.copytree(base_data_dir, data_dir)
+
+
     def init(self, allows_streaming=False, initdb_params=[]):
         """ Performs initdb """
 
@@ -258,19 +286,13 @@ def init(self, allows_streaming=False, initdb_params=[]):
             return self
 
         # initialize cluster
-        os.makedirs(self.data_dir)
-        initdb = self.get_bin_path("initdb")
-        with open(os.path.join(self.logs_dir, "initdb.log"), "a") as file_out:
-            ret = subprocess.call(
-                [initdb, self.data_dir, "-N"] + initdb_params,
-                stdout=file_out,
-                stderr=subprocess.STDOUT
-            )
-            if ret:
-                raise ClusterException("Cluster initialization failed")
+        if initdb_params:
+            self.initdb(self.data_dir, initdb_params)
+        else:
+            self._setup_data_dir(self.data_dir)
 
         # add parameters to config file
-        with open(postgres_conf, "a") as conf:
+        with open(postgres_conf, "w") as conf:
             conf.write(
                 "fsync = off\n"
                 "log_statement = all\n"
@@ -286,7 +308,6 @@ def init(self, allows_streaming=False, initdb_params=[]):
                     "max_wal_senders = 5\n"
                     "wal_keep_segments = 20\n"
                     "max_wal_size = 128MB\n"
-                    "shared_buffers = 1MB\n"
                     "wal_log_hints = on\n"
                     "hot_standby = on\n"
                     "max_connections = 10\n")
@@ -369,7 +390,9 @@ def pg_ctl(self, command, params={}, command_options=[]):
 
             if res > 0:
                 with open(self.error_filename, "r") as errfile:
-                    raise ClusterException(errfile.readlines()[-1])
+                    text = errfile.readlines()[-1]
+                    text += 'Logs at: %s' % self.logs_dir
+                    raise ClusterException(text)
 
     def start(self, params={}):
         """ Starts cluster """
@@ -434,7 +457,7 @@ def get_control_data(self):
             lines = subprocess.check_output(
                 [pg_controldata] + ["-D", self.data_dir],
                 stderr=subprocess.STDOUT
-            ).splitlines()
+            ).decode("utf-8").splitlines()
         except subprocess.CalledProcessError as e:
             raise PgcontroldataException(e.output, e.cmd)
 
@@ -489,8 +512,7 @@ def cleanup(self):
                 pass
 
         # remove data directory
-        shutil.rmtree(self.data_dir)
-
+        shutil.rmtree(self.data_dir, ignore_errors=True)
         return self
 
     def psql(self, dbname, query=None, filename=None, username=None):
@@ -597,7 +619,7 @@ def backup(self, name):
         backup_path = os.path.join(self.base_dir, name)
         os.makedirs(backup_path)
         params = [pg_basebackup, "-D", backup_path, "-p {}".format(
-            self.port), "-x"]
+            self.port), "-X", "fetch"]
         with open(self.output_filename, "a") as file_out, \
                 open(self.error_filename, "a") as file_err:
             ret = subprocess.call(
@@ -726,7 +748,16 @@ def clean_all():
     global registered_nodes
     for node in registered_nodes:
         node.cleanup()
+
+    for cat in tmp_dirs:
+        if os.path.exists():
+            shutil.rmtree(cat, ignore_errors=True)
+
+        if cat == base_data_dir:
+            base_data_dir = None
+
     registered_nodes = []
+    tmp_dirs = []
 
 
 def stop_all():

testgres/tests/test_simple.py

@@ -6,13 +6,13 @@
 import tempfile
 import logging.config
 
-from testgres import get_new_node, stop_all
+from testgres import get_new_node, stop_all, get_config, clean_all
 
 
 class SimpleTest(unittest.TestCase):
 
     def teardown(self):
-        # clean_all()
+        clean_all()
         stop_all()
 
     @unittest.skip("demo")
@@ -41,15 +41,21 @@ def test_backup_and_replication(self):
         self.assertEqual(len(res), 1)
         self.assertEqual(res[0], (1, 2))
 
+        # Prepare the query which would check whether record reached replica
+        # (It is slightly different for Postgres 9.6 and Postgres 10+)
+        if get_config()['VERSION_NUM'] >= 1000000:
+            wait_lsn = 'SELECT pg_current_wal_lsn() <= replay_lsn '           \
+                'FROM pg_stat_replication WHERE application_name = \'%s\''    \
+                % replica.name
+        else:
+            wait_lsn = 'SELECT pg_current_xlog_location() <= replay_location '\
+                'FROM pg_stat_replication WHERE application_name = \'%s\''    \
+                % replica.name
+
         # Insert into master node
         node.psql('postgres', 'insert into abc values (3, 4)')
         # Wait until data syncronizes
-        node.poll_query_until(
-            'postgres',
-            'SELECT pg_current_xlog_location() <= replay_location '
-            'FROM pg_stat_replication WHERE application_name = \'%s\''
-            % replica.name)
-        # time.sleep(0.5)
+        node.poll_query_until('postgres', wait_lsn)
         # Check that this record was exported to replica
         res = replica.execute('postgres', 'select * from abc')
         self.assertEqual(len(res), 2)
@@ -81,9 +87,10 @@ def test_users(self):
         node.psql('postgres', 'create role test_user login')
         value = node.safe_psql('postgres', 'select 1', username='test_user')
         self.assertEqual(value, six.b('1\n'))
+        node.stop()
 
     def test_logging(self):
-        regex = re.compile('\w+:\s{1}LOG:.*')
+        regex = re.compile('.+?LOG:.*')
         logfile = tempfile.NamedTemporaryFile('w', delete=True)
 
         log_conf = {