Fixes and flake8 cleanups

Fenimorkin · Fenimorkin · commit 458e2b0891d6 · 2018-04-28T18:22:57.000+03:00
diff --git a/testgres/consts.py b/testgres/consts.py
@@ -25,10 +25,6 @@
 UTILS_LOG_FILE = "utils.log"
 BACKUP_LOG_FILE = "backup.log"
 
-# name of dump/restore formates
-DUMP_DIRECTORY = "directory"
-DEFAULT_DUMP_FORMAT = "plain"
-
 # defaults for node settings
 MAX_REPLICATION_SLOTS = 10
 MAX_WAL_SENDERS = 10
diff --git a/testgres/enums.py b/testgres/enums.py
@@ -85,3 +85,14 @@ def from_process(process):
 
         # default
         return ProcessType.Unknown
+
+
+class DumpFormat(Enum):
+    """
+    Available dump formats
+    """
+
+    Plain = 'plain'
+    Custom = 'custom'
+    Directory = 'directory'
+    Tar = 'tar'
diff --git a/testgres/node.py b/testgres/node.py
@@ -10,7 +10,10 @@
 from six import raise_from, iteritems
 from tempfile import mkstemp, mkdtemp
 
-from .enums import NodeStatus, ProcessType
+from .enums import \
+    NodeStatus, \
+    ProcessType, \
+    DumpFormat
 
 from .cache import cached_initdb
 
@@ -32,8 +35,6 @@
     RECOVERY_CONF_FILE, \
     PG_LOG_FILE, \
     UTILS_LOG_FILE, \
-    DEFAULT_DUMP_FORMAT, \
-    DUMP_DIRECTORY, \
     PG_PID_FILE
 
 from .consts import \
@@ -805,23 +806,27 @@ def safe_psql(self, query=None, **kwargs):
 
         return out
 
-    def dump(self, filename=None, dbname=None, username=None, format=DEFAULT_DUMP_FORMAT):
+    def dump(self,
+             filename=None,
+             dbname=None,
+             username=None,
+             format=DumpFormat.Plain):
         """
         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.
-            format: format argument plain/custom/directory/tar
+            format: format argument plain/custom/directory/tar.
 
         Returns:
             Path to a file containing dump.
         """
 
         def tmpfile():
-            if format == DUMP_DIRECTORY:
+            if format == DumpFormat.Directory:
                 fname = mkdtemp(prefix=TMP_DUMP)
             else:
                 fd, fname = mkstemp(prefix=TMP_DUMP)
@@ -852,7 +857,7 @@ def restore(self, filename, dbname=None, username=None):
         Restore database from pg_dump's file.
 
         Args:
-            filename: input file.
+            filename: database dump taken by pg_dump in custom/directory/tar formats.
             dbname: database name to connect to.
             username: database user name.
         """
@@ -862,17 +867,13 @@ def restore(self, filename, dbname=None, username=None):
         username = username or default_username()
 
         _params = [
-            get_bin_path("pg_restore"),
-            "-p", str(self.port),
-            "-h", self.host,
-            "-U", username,
-            "-d", dbname,
+            get_bin_path("pg_restore"), "-p",
+            str(self.port), "-h", self.host, "-U", username, "-d", dbname,
             filename
         ]
 
         execute_utility(_params, self.utils_log_name)
 
-
     @method_decorator(positional_args_hack(['dbname', 'query']))
     def poll_query_until(self,
                          query,
diff --git a/tests/test_simple.py b/tests/test_simple.py
@@ -10,7 +10,6 @@
 import unittest
 
 import logging.config
-import shutil
 
 from contextlib import contextmanager
 from distutils.version import LooseVersion
