replace pg_version_ge() with explicit per-node version

funbringer · funbringer · commit 72d03735dace · 2018-06-01T16:54:12.000+03:00
diff --git a/testgres/node.py b/testgres/node.py
@@ -66,10 +66,11 @@
 from .pubsub import Publication, Subscription
 
 from .utils import \
+    PgVer, \
     eprint, \
     get_bin_path, \
+    get_pg_version, \
     file_tail, \
-    pg_version_ge, \
     reserve_port, \
     release_port, \
     execute_utility, \
@@ -113,6 +114,7 @@ def __init__(self, name=None, port=None, base_dir=None):
         """
 
         # private
+        self._pg_version = PgVer(get_pg_version())
         self._should_free_port = port is None
         self._base_dir = base_dir
         self._logger = None
@@ -489,7 +491,7 @@ def get_auth_method(t):
             if allow_streaming:
 
                 # select a proper wal_level for PostgreSQL
-                if pg_version_ge('9.6'):
+                if self._pg_version >= '9.6':
                     wal_level = "replica"
                 else:
                     wal_level = "hot_standby"
@@ -504,7 +506,7 @@ def get_auth_method(t):
                                                       wal_level))  # yapf: disable
 
             if allow_logical:
-                if not pg_version_ge('10'):
+                if self._pg_version < '10':
                     raise InitNodeException(
                         "Logical replication is only available for Postgres 10 "
                         "and newer")
@@ -568,7 +570,7 @@ def get_control_data(self):
 
         # this one is tricky (blame PG 9.4)
         _params = [get_bin_path("pg_controldata")]
-        _params += ["-D"] if pg_version_ge('9.5') else []
+        _params += ["-D"] if self._pg_version >= '9.5' else []
         _params += [self.data_dir]
 
         data = execute_utility(_params, self.utils_log_file)
@@ -1042,7 +1044,7 @@ def catchup(self, dbname=None, username=None):
         if not self.master:
             raise TestgresException("Node doesn't have a master")
 
-        if pg_version_ge('10'):
+        if self._pg_version >= '10':
             poll_lsn = "select pg_catalog.pg_current_wal_lsn()::text"
             wait_lsn = "select pg_catalog.pg_last_wal_replay_lsn() >= '{}'::pg_lsn"
         else:
diff --git a/testgres/utils.py b/testgres/utils.py
@@ -22,6 +22,9 @@
 # ports used by nodes
 bound_ports = set()
 
+# re-export version type
+PgVer = LooseVersion
+
 
 def reserve_port():
     """
@@ -180,17 +183,6 @@ def get_pg_version():
     return version
 
 
-def pg_version_ge(version):
-    """
-    Check if PostgreSQL is 'version' or newer.
-    """
-
-    cur_ver = LooseVersion(get_pg_version())
-    min_ver = LooseVersion(version)
-
-    return cur_ver >= min_ver
-
-
 def file_tail(f, num_lines):
     """
     Get last N lines of a file.
diff --git a/tests/test_simple.py b/tests/test_simple.py
@@ -12,7 +12,6 @@
 import logging.config
 
 from contextlib import contextmanager
-from distutils.version import LooseVersion
 from shutil import rmtree
 
 from testgres import \
@@ -32,16 +31,24 @@
 
 from testgres import \
     NodeStatus, \
+    ProcessType, \
     IsolationLevel, \
     get_new_node
 
 from testgres import \
     get_bin_path, \
-    get_pg_config
+    get_pg_config, \
+    get_pg_version
 
+# NOTE: those are ugly imports
 from testgres import bound_ports
-from testgres.utils import pg_version_ge
-from testgres.enums import ProcessType
+from testgres.utils import PgVer
+
+
+def pg_version_ge(version):
+    cur_ver = PgVer(get_pg_version())
+    min_ver = PgVer(version)
+    return cur_ver >= min_ver
 
 
 def util_exists(util):
@@ -110,7 +117,7 @@ def test_init_after_cleanup(self):
             node.init().start().execute('select 1')
 
     @unittest.skipUnless(util_exists('pg_resetwal'), 'might be missing')
-    @unittest.skipUnless(pg_version_ge('9.6'), 'query works on 9.6+')
+    @unittest.skipUnless(pg_version_ge('9.6'), 'requires 9.6+')
     def test_init_unique_system_id(self):
         # this function exists in PostgreSQL 9.6+
         query = 'select system_identifier from pg_control_system()'
@@ -828,9 +835,9 @@ def test_exceptions(self):
         str(QueryException('msg', 'query'))
 
     def test_version_management(self):
-        a = LooseVersion('10.0')
-        b = LooseVersion('10')
-        c = LooseVersion('9.6.5')
+        a = PgVer('10.0')
+        b = PgVer('10')
+        c = PgVer('9.6.5')
 
         self.assertTrue(a > b)
         self.assertTrue(b > c)
