Merge pull request #33 from Valeria1235/simplified_enums

funbringer · web-flow · commit 3175125420d0 · 2018-03-04T00:05:07.000+03:00
Simplified enums
diff --git a/testgres/__init__.py b/testgres/__init__.py
@@ -9,13 +9,13 @@
     pop_config
 
 from .connection import \
-    IsolationLevel, \
     NodeConnection, \
     InternalError, \
     ProgrammingError
 
 from .exceptions import *
-from .node import NodeStatus, PostgresNode
+from .enums import *
+from .node import PostgresNode
 
 from .utils import \
     reserve_port, \
diff --git a/testgres/backup.py b/testgres/backup.py
@@ -6,6 +6,8 @@
 from six import raise_from
 from tempfile import mkdtemp
 
+from .enums import XLogMethod
+
 from .consts import \
     DATA_DIR, \
     TMP_NODE, \
@@ -50,7 +52,13 @@ def __init__(self,
         if not node.status():
             raise BackupException('Node must be running')
 
-        # yapf: disable
+        # Check arguments
+        if not isinstance(xlog_method, XLogMethod):
+            try:
+                xlog_method = XLogMethod(xlog_method)
+            except ValueError:
+                raise BackupException('Invalid xlog_method "{}"'.format(xlog_method))
+
         # Set default arguments
         username = username or default_username()
         base_dir = base_dir or mkdtemp(prefix=TMP_BACKUP)
@@ -72,7 +80,7 @@ def __init__(self,
             "-h", node.host,
             "-U", username,
             "-D", data_dir,
-            "-X", xlog_method
+            "-X", xlog_method.value
         ]
         execute_utility(_params, self.log_file)
 
diff --git a/testgres/connection.py b/testgres/connection.py
@@ -9,12 +9,12 @@
     except ImportError:
         raise ImportError("You must have psycopg2 or pg8000 modules installed")
 
+from .enums import IsolationLevel
+
 from .defaults import \
     default_dbname, \
     default_username
 
-from enum import Enum
-
 from .exceptions import QueryException
 
 
@@ -23,14 +23,6 @@
 ProgrammingError = pglib.ProgrammingError
 
 
-class IsolationLevel(Enum):
-    """
-    Transaction isolation level for NodeConnection
-    """
-
-    ReadUncommitted, ReadCommitted, RepeatableRead, Serializable = range(4)
-
-
 class NodeConnection(object):
     """
     Transaction wrapper returned by Node
@@ -72,39 +64,21 @@ def __exit__(self, type, value, traceback):
         self.close()
 
     def begin(self, isolation_level=IsolationLevel.ReadCommitted):
-        # yapf: disable
-        levels = [
-            'read uncommitted',
-            'read committed',
-            'repeatable read',
-            'serializable'
-        ]
-
-        # Check if level is an IsolationLevel
-        if (isinstance(isolation_level, IsolationLevel)):
-
-            # Get index of isolation level
-            level_idx = isolation_level.value
-            assert level_idx in range(4)
-
-            # Replace isolation level with its name
-            isolation_level = levels[level_idx]
-
-        else:
+        # Check if level isn't an IsolationLevel
+        if not isinstance(isolation_level, IsolationLevel):
             # Get name of isolation level
             level_str = str(isolation_level).lower()
 
             # Validate level string
-            if level_str not in levels:
+            try:
+                isolation_level = IsolationLevel(level_str)
+            except ValueError:
                 error = 'Invalid isolation level "{}"'
                 raise QueryException(error.format(level_str))
 
-            # Replace isolation level with its name
-            isolation_level = level_str
-
         # Set isolation level
         cmd = 'SET TRANSACTION ISOLATION LEVEL {}'
-        self.cursor.execute(cmd.format(isolation_level))
+        self.cursor.execute(cmd.format(isolation_level.value))
 
         return self
 
diff --git a/testgres/consts.py b/testgres/consts.py
@@ -1,5 +1,7 @@
 # coding: utf-8
 
+from .enums import XLogMethod
+
 # names for dirs in base_dir
 DATA_DIR = "data"
 LOGS_DIR = "logs"
@@ -25,4 +27,4 @@
 BACKUP_LOG_FILE = "backup.log"
 
 # default argument value
-DEFAULT_XLOG_METHOD = "fetch"
+DEFAULT_XLOG_METHOD = XLogMethod.fetch
diff --git a/testgres/enums.py b/testgres/enums.py
@@ -0,0 +1,37 @@
+from enum import Enum, IntEnum
+
+
+class XLogMethod(Enum):
+    """
+    Available WAL methods for NodeBackup
+    """
+
+    none = 'none'
+    fetch = 'fetch'
+    stream = 'stream'
+
+
+class IsolationLevel(Enum):
+    """
+    Transaction isolation level for NodeConnection
+    """
+
+    ReadUncommitted = 'read uncommitted'
+    ReadCommitted = 'read committed'
+    RepeatableRead = 'repeatable read'
+    Serializable = 'serializable'
+
+
+class NodeStatus(IntEnum):
+    """
+    Status of a PostgresNode
+    """
+
+    Running, Stopped, Uninitialized = range(3)
+
+    # for Python 3.x
+    def __bool__(self):
+        return self == NodeStatus.Running
+
+    # for Python 2.x
+    __nonzero__ = __bool__
diff --git a/testgres/node.py b/testgres/node.py
@@ -6,11 +6,12 @@
 import subprocess
 import time
 
-from enum import Enum
 from shutil import rmtree
 from six import raise_from
 from tempfile import mkstemp, mkdtemp
 
+from .enums import NodeStatus
+
 from .cache import cached_initdb
 
 from .config import testgres_config
@@ -59,20 +60,7 @@
     release_port, \
     execute_utility
 
-
-class NodeStatus(Enum):
-    """
-    Status of a PostgresNode
-    """
-
-    Running, Stopped, Uninitialized = range(3)
-
-    # for Python 3.x
-    def __bool__(self):
-        return self.value == NodeStatus.Running.value
-
-    # for Python 2.x
-    __nonzero__ = __bool__
+from .backup import NodeBackup
 
 
 class PostgresNode(object):
@@ -858,7 +846,6 @@ def backup(self, **kwargs):
             A smart object of type NodeBackup.
         """
 
-        from .backup import NodeBackup
         return NodeBackup(node=self, **kwargs)
 
     def replicate(self, name=None, **kwargs):
diff --git a/tests/test_simple.py b/tests/test_simple.py
@@ -360,6 +360,13 @@ def test_backup_exhaust(self):
                 with self.assertRaises(BackupException):
                     backup.spawn_primary()
 
+    def test_backup_wrong_xlog_method(self):
+        with get_new_node() as node:
+            node.init(allow_streaming=True).start()
+
+            with self.assertRaises(BackupException, msg='Invalid xlog_method "wrong"'):
+                node.backup(xlog_method='wrong')
+
     def test_replicate(self):
         with get_new_node() as node:
             node.init(allow_streaming=True).start()
