Fix tests for pg_dump and pg_restore

postgrespro · funbringer · May 8, 2018 · Feb 4, 2018 · Apr 28, 2018 · Apr 28, 2018
commit e76232f82a2015cb1b3ecb97e979da14e0a405fe
diff --git a/testgres/consts.py b/testgres/consts.py
@@ -13,5 +13,6 @@
 UTILS_LOG_FILE = "utils.log"
 BACKUP_LOG_FILE = "backup.log"
 
-# default argument value
+# default argument values
 DEFAULT_XLOG_METHOD = "fetch"
+DEFAULT_DUMP_FORMAT = "plain"
diff --git a/testgres/node.py b/testgres/node.py
@@ -28,7 +28,8 @@
     RECOVERY_CONF_FILE, \
     PG_LOG_FILE, \
     UTILS_LOG_FILE, \
-    DEFAULT_XLOG_METHOD
+    DEFAULT_XLOG_METHOD, \
+    DEFAULT_DUMP_FORMAT
 
 from .exceptions import \
     CatchUpException,   \
@@ -606,8 +607,9 @@ def cleanup(self, max_attempts=3):
     @method_decorator(positional_args_hack(['query'], ['dbname', 'query']))
     def psql(self,
              query=None,
-             username=None,
              filename=None,
+             dbname=None,
+             username=None,
              input=None):
         """
         Execute a query using psql.
@@ -687,32 +689,34 @@ def safe_psql(self,
 
         return out
 
-    def dump(self, dbname, username=None, filename=None, formate=None):
+    def dump(self, filename=None, dbname=None, username=None, format=DEFAULT_DUMP_FORMAT):
         """
         Dump database into a file using pg_dump.
         NOTE: the file is not removed automatically.
 
         Args:
-            filename: database dump taken by pg_dump.
             dbname: database name to connect to.
             username: database user name.
             filename: output file.
-            formate: format argument p/c/d/t
+            format: format argument plain/custom/directory/tar
 
         Returns:
             Path to a file containing dump.
         """
 
         def tmpfile():
-            fd, fname = tempfile.mkstemp()
-            os.close(fd)
+            if format == 'directory':
+                fname = tempfile.mkdtemp()
+            else:
+                fd, fname = tempfile.mkstemp()
+                os.close(fd)
             return fname
 
         # Set default arguments
-        username = username or _default_username()
-        formate = formate or "p"
-        f, filename = filename or tempfile.mkstemp()
-        os.close(f)
+        dbname = dbname or default_dbname()
+        username = username or default_username()
+        filename = filename or tmpfile()
+
 
         # yapf: disable
         _params = [
@@ -722,7 +726,7 @@ def tmpfile():
             "-f", filename,
             "-U", username,
             "-d", dbname,
-            "-F", formate
+            "-F", format
         ]
 
         execute_utility(_params, self.utils_log_name)
@@ -734,14 +738,40 @@ def restore(self, filename, dbname=None, username=None):
         Restore database from pg_dump's file.
 
         Args:
-            dbname: database name to connect to.
-            username: database user name.
-            filename: database dump taken by pg_dump.
+            filename: input file.
             dbname: database name to connect to.
             username: database user name.
         """
 
-        self.psql(dbname=dbname, username=username, filename=filename)
+        # self.psql(filename=filename, dbname=dbname, username=username)
+                # yapf: disable
+        # _params = [
+        #     get_bin_path("pg_restore"),
+        #     "-p", str(self.port),
+        #     "-h", self.host,
+        #           filename,
+        #     "-U", username,
+        #     "-d", dbname
+        # ]
+
+        # Set default arguments
+        dbname = dbname or default_dbname()
+        username = username or default_username()
+
+        _params = [
+            get_bin_path("pg_restore"),
+            "-p", str(self.port),
+            "-h", self.host,
+            #       filename,
+            # "-c",
+            # "-U", username,
+            "-d", dbname,
+            filename
+        ]
+        print(_params)
+
+        execute_utility(_params, self.utils_log_name)
+
 
     @method_decorator(positional_args_hack(['dbname', 'query']))
     def poll_query_until(self,

diff --git a/tests/test_pgdump.py b/tests/test_pgdump.py
diff --git a/tests/test_simple.py b/tests/test_simple.py
@@ -9,6 +9,9 @@
 import unittest
 
 import logging.config
+import shutil
+
+import pdb
 
 from distutils.version import LooseVersion
 
@@ -368,19 +371,101 @@ def test_dump(self):
 
             node1.execute(query_create)
 
-            # take a new dump
+            # take a new dump plain format
             dump = node1.dump()
             self.assertTrue(os.path.isfile(dump))
+            with get_new_node().init().start() as node2:
+                # restore dump
+                node2.restore(filename=dump)
+                res = node2.execute(query_select)
+                self.assertListEqual(res, [(1, ), (2, )])
+            # finally, remove dump
+            os.remove(dump)
+
+    def test_dump_plain(self):
+        query_create = 'create table test as select generate_series(1, 2) as val'
+        query_select = 'select * from test order by val asc'
 
+        with get_new_node().init().start() as node1:
+
+            node1.execute(query_create)
+
+            # take a new dump plain format
+            dump = node1.dump(format='plain')
+            self.assertTrue(os.path.isfile(dump))
             with get_new_node().init().start() as node2:
+                # node2.restore(filename=dump)
+                node2.psql(filename=dump, dbname=None, username=None) #add
+                res = node2.execute(query_select)
+                self.assertListEqual(res, [(1, ), (2, )])
+            os.remove(dump)
 
-                # restore dump
+            # with get_new_node().init().start() as node2:
+            #     for f in ['plain', 'custom', 'directory', 'tar']:
+            #         # dump
+            #         dump = node1.dump(format=f)
+            #         # create database
+            #         # connect
+            #         # restore
+            #         node2.restore(filename=f)
+            #         # check
+            #         res = node2.execute(query_select)
+            #         self.assertListEqual(res, [(1, ), (2, )])
+            #         # drop database
+            #         os.remove(f)
+
+    def test_dump_custom(self):
+        query_create = 'create table test as select generate_series(1, 2) as val'
+        query_select = 'select * from test order by val asc'
+
+        with get_new_node().init().start() as node1:
+
+            node1.execute(query_create)
+
+            # take a new dump custom format
+            dump = node1.dump(format='custom')
+            self.assertTrue(os.path.isfile(dump))
+            with get_new_node().init().start() as node2:
                 node2.restore(filename=dump)
+                res = node2.execute(query_select)
+                print(res)
+                self.assertListEqual(res, [(1, ), (2, )])
+            os.remove(dump)
+
+    def test_dump_directory(self):
+        query_create = 'create table test as select generate_series(1, 2) as val'
+        query_select = 'select * from test order by val asc'
+
+        with get_new_node().init().start() as node1:
 
+            node1.execute(query_create)
+
+            # take a new dump directory format
+            dump = node1.dump(format='directory')
+            self.assertTrue(os.path.isdir(dump))
+            with get_new_node().init().start() as node2:
+                node2.restore(filename=dump)
                 res = node2.execute(query_select)
                 self.assertListEqual(res, [(1, ), (2, )])
+            shutil.rmtree(dump, ignore_errors=True)
 
-            # finally, remove dump
+    def test_dump_tar(self):
+        query_create = 'create table test as select generate_series(1, 2) as val'
+        query_select = 'select * from test order by val asc'
+
+        with get_new_node().init().start() as node1:
+
+            node1.execute(query_create)
+
+            # take a new dump tar format
+            dump = node1.dump(format='tar')
+            self.assertTrue(os.path.isfile(dump))
+            with get_new_node().init().start() as node2:
+                node2.restore(filename=dump)
+                print("Restore finished")
+                pdb.set_trace()
+                res = node2.execute(query_select)
+                self.assertListEqual(res, [(1, ), (2, )])
             os.remove(dump)
 
     def test_users(self):