PBCKP-152

Author: v.shepard

testgres/backup.py

@@ -2,7 +2,6 @@
 
 import os
 
-from shutil import rmtree, copytree
 from six import raise_from
 
 from .enums import XLogMethod
@@ -14,8 +13,6 @@
     PG_CONF_FILE, \
     BACKUP_LOG_FILE
 
-from .defaults import default_username
-
 from .exceptions import BackupException
 
 from .utils import \
@@ -80,7 +77,7 @@ def __init__(self,
             "-D", data_dir,
             "-X", xlog_method.value
         ]  # yapf: disable
-        execute_utility(_params, self.log_file, hostname=node.hostname, ssh_key=node.ssh_key)
+        execute_utility(_params, self.log_file, self.os_ops)
 
     def __enter__(self):
         return self
@@ -113,7 +110,7 @@ def _prepare_dir(self, destroy):
 
             try:
                 # Copy backup to new data dir
-                copytree(data1, data2)
+                self.os_ops.copytree(data1, data2)
             except Exception as e:
                 raise_from(BackupException('Failed to copy files'), e)
         else:

testgres/cache.py

@@ -4,7 +4,8 @@
 
 from six import raise_from
 
-from .os_ops import OsOperations
+from .op_ops.local_ops import LocalOperations
+from .op_ops.os_ops import OsOperations
 from .config import testgres_config
 
 from .consts import XLOG_CONTROL_FILE
@@ -20,16 +21,15 @@
     execute_utility
 
 
-def cached_initdb(data_dir, logfile=None, hostname='localhost', ssh_key=None, params=None):
+def cached_initdb(data_dir, logfile=None, params=None, os_ops: OsOperations = LocalOperations()):
     """
     Perform initdb or use cached node files.
     """
-    os_ops = OsOperations(hostname=hostname, ssh_key=ssh_key)
 
     def call_initdb(initdb_dir, log=None):
         try:
             _params = [get_bin_path("initdb"), "-D", initdb_dir, "-N"]
-            execute_utility(_params + (params or []), log, hostname=hostname, ssh_key=ssh_key)
+            execute_utility(_params + (params or []), log, os_ops)
         except ExecUtilException as e:
             raise_from(InitNodeException("Failed to run initdb"), e)
 
@@ -42,7 +42,7 @@ def call_initdb(initdb_dir, log=None):
         # Initialize cached initdb
 
         if not os_ops.path_exists(cached_data_dir) or \
-           not os_ops.listdir(cached_data_dir):
+                not os_ops.listdir(cached_data_dir):
             call_initdb(cached_data_dir)
 
         try:
@@ -60,7 +60,7 @@ def call_initdb(initdb_dir, log=None):
 
                 # XXX: build new WAL segment with our system id
                 _params = [get_bin_path("pg_resetwal"), "-D", data_dir, "-f"]
-                execute_utility(_params, logfile, hostname=hostname, ssh_key=ssh_key)
+                execute_utility(_params, logfile, os_ops)
 
         except ExecUtilException as e:
             msg = "Failed to reset WAL for system id"

testgres/config.py

@@ -6,8 +6,8 @@
 
 from contextlib import contextmanager
 from shutil import rmtree
-from tempfile import mkdtemp
 
+from .op_ops.local_ops import LocalOperations
 from .consts import TMP_CACHE
 
 
@@ -137,12 +137,9 @@ def copy(self):
 
 
 @atexit.register
-def _rm_cached_initdb_dirs(os_ops=None):
+def _rm_cached_initdb_dirs(os_ops=LocalOperations()):
     for d in cached_initdb_dirs:
-        if os_ops:
-            os_ops.rmtree(d, ignore_errors=True)
-        else:
-            rmtree(d, ignore_errors=True)
+        os_ops.rmdirs(d, ignore_errors=True)
 
 
 def push_config(**options):
@@ -205,4 +202,4 @@ def configure_testgres(**options):
 
 
 # NOTE: assign initial cached dir for initdb
-testgres_config.cached_initdb_dir = mkdtemp(prefix=TMP_CACHE)
+testgres_config.cached_initdb_dir = testgres_config.os_ops.mkdtemp(prefix=TMP_CACHE)

testgres/connection.py

@@ -37,7 +37,7 @@ def __init__(self,
 
         # Set default arguments
         dbname = dbname or default_dbname()
-        username = username or default_username()
+        username = username or default_username(node.os_ops)
 
         self._node = node
 

testgres/defaults.py

@@ -1,9 +1,10 @@
 import datetime
-import getpass
 import os
 import struct
 import uuid
 
+from .op_ops.local_ops import LocalOperations
+
 
 def default_dbname():
     """
@@ -13,15 +14,11 @@ def default_dbname():
     return 'postgres'
 
 
-def default_username(os_ops=None):
+def default_username(os_ops=LocalOperations()):
     """
     Return default username (current user).
     """
-    if os_ops:
-        user = os_ops.get_user()
-    else:
-        user = getpass.getuser()
-    return user
+    return os_ops.get_user()
 
 
 def generate_app_name():
@@ -32,7 +29,7 @@ def generate_app_name():
     return 'testgres-{}'.format(str(uuid.uuid4()))
 
 
-def generate_system_id(os_ops=None):
+def generate_system_id(os_ops=LocalOperations()):
     """
     Generate a new 64-bit unique system identifier for node.
     """
@@ -47,10 +44,7 @@ def generate_system_id(os_ops=None):
     system_id = 0
     system_id |= (secs << 32)
     system_id |= (usecs << 12)
-    if os_ops:
-        system_id |= (os_ops.get_pid() & 0xFFF)
-    else:
-        system_id |= (os.getpid() & 0xFFF)
+    system_id |= (os_ops.get_pid() & 0xFFF)
 
     # pack ULL in native byte order
     return struct.pack('=Q', system_id)