small refactoring

Author: funbringer

testgres/backup.py

@@ -1,13 +1,11 @@
 # coding: utf-8
 
 import os
-import shutil
 
+from shutil import rmtree, copytree
 from six import raise_from
 from tempfile import mkdtemp
 
-from .config import testgres_config
-
 from .consts import \
     DATA_DIR, \
     TMP_NODE, \
@@ -55,8 +53,7 @@ def __init__(self,
         # yapf: disable
         # Set default arguments
         username = username or default_username()
-        base_dir = base_dir or mkdtemp(prefix=TMP_BACKUP,
-                                       dir=testgres_config.temp_dir)
+        base_dir = base_dir or mkdtemp(prefix=TMP_BACKUP)
 
         # public
         self.original_node = node
@@ -103,15 +100,14 @@ def _prepare_dir(self, destroy):
         available = not destroy
 
         if available:
-            dest_base_dir = mkdtemp(prefix=TMP_NODE,
-                                    dir=testgres_config.temp_dir)
+            dest_base_dir = mkdtemp(prefix=TMP_NODE)
 
             data1 = os.path.join(self.base_dir, DATA_DIR)
             data2 = os.path.join(dest_base_dir, DATA_DIR)
 
             try:
                 # Copy backup to new data dir
-                shutil.copytree(data1, data2)
+                copytree(data1, data2)
             except Exception as e:
                 raise_from(BackupException('Failed to copy files'), e)
         else:
@@ -140,8 +136,7 @@ def spawn_primary(self, name=None, destroy=True):
 
         # Build a new PostgresNode
         from .node import PostgresNode
-        node = PostgresNode(name=name,
-                            base_dir=base_dir)
+        node = PostgresNode(name=name, base_dir=base_dir)
 
         # New nodes should always remove dir tree
         node._should_rm_dirs = True
@@ -164,8 +159,7 @@ def spawn_replica(self, name=None, destroy=True):
         """
 
         # Build a new PostgresNode
-        node = self.spawn_primary(name=name,
-                                  destroy=destroy)
+        node = self.spawn_primary(name=name, destroy=destroy)
 
         # Assign it a master and a recovery file (private magic)
         node._assign_master(self.original_node)
@@ -175,5 +169,5 @@ def spawn_replica(self, name=None, destroy=True):
 
     def cleanup(self):
         if self._available:
-            shutil.rmtree(self.base_dir, ignore_errors=True)
+            rmtree(self.base_dir, ignore_errors=True)
             self._available = False

testgres/cache.py

@@ -2,8 +2,8 @@
 
 import io
 import os
-import shutil
 
+from shutil import copytree
 from six import raise_from
 
 from .config import testgres_config
@@ -46,7 +46,7 @@ def call_initdb(initdb_dir, log=None):
 
         try:
             # Copy cached initdb to current data dir
-            shutil.copytree(cached_data_dir, data_dir)
+            copytree(cached_data_dir, data_dir)
 
             # Assign this node a unique system id if asked to
             if testgres_config.cached_initdb_unique:

testgres/config.py

@@ -2,9 +2,10 @@
 
 import atexit
 import copy
-import shutil
+import tempfile
 
 from contextlib import contextmanager
+from shutil import rmtree
 from tempfile import mkdtemp
 
 from .consts import TMP_CACHE
@@ -39,8 +40,6 @@ class GlobalConfig(object):
 
     cache_pg_config = True
 
-    temp_dir = None
-
     use_python_logging = False
     error_log_lines = 20
 
@@ -59,6 +58,14 @@ def cached_initdb_dir(self, value):
         if value:
             cached_initdb_dirs.add(value)
 
+    @property
+    def temp_dir(self):
+        return tempfile.tempdir
+
+    @temp_dir.setter
+    def temp_dir(self, value):
+        tempfile.tempdir = value
+
     def __init__(self, **options):
         self.update(options)
 
@@ -110,9 +117,9 @@ def copy(self):
 
 
 @atexit.register
-def rm_cached_initdb_dirs():
+def _rm_cached_initdb_dirs():
     for d in cached_initdb_dirs:
-        shutil.rmtree(d, ignore_errors=True)
+        rmtree(d, ignore_errors=True)
 
 
 def push_config(**options):
@@ -170,7 +177,5 @@ def configure_testgres(**options):
     testgres_config.update(options)
 
 
-# yapf: disable
 # NOTE: assign initial cached dir for initdb
-testgres_config.cached_initdb_dir = mkdtemp(prefix=TMP_CACHE,
-                                            dir=testgres_config.temp_dir)
+testgres_config.cached_initdb_dir = mkdtemp(prefix=TMP_CACHE)

testgres/node.py

@@ -2,12 +2,12 @@
 
 import io
 import os
-import shutil
 import six
 import subprocess
 import time
 
 from enum import Enum
+from shutil import rmtree
 from six import raise_from
 from tempfile import mkstemp, mkdtemp
 
@@ -203,8 +203,7 @@ def _create_recovery_conf(self, username):
 
     def _prepare_dirs(self):
         if not self.base_dir:
-            self.base_dir = mkdtemp(prefix=TMP_NODE,
-                                    dir=testgres_config.temp_dir)
+            self.base_dir = mkdtemp(prefix=TMP_NODE)
 
         if not os.path.exists(self.base_dir):
             os.makedirs(self.base_dir)
@@ -613,7 +612,7 @@ def cleanup(self, max_attempts=3):
         else:
             rm_dir = self.data_dir    # just data, save logs
 
-        shutil.rmtree(rm_dir, ignore_errors=True)
+        rmtree(rm_dir, ignore_errors=True)
 
         return self
 
@@ -711,7 +710,7 @@ def dump(self, filename=None, dbname=None, username=None):
         """
 
         def tmpfile():
-            fd, fname = mkstemp(prefix=TMP_DUMP, dir=testgres_config.temp_dir)
+            fd, fname = mkstemp(prefix=TMP_DUMP)
             os.close(fd)
             return fname
 

tests/test_simple.py

@@ -2,7 +2,6 @@
 # coding: utf-8
 
 import os
-import shutil
 import subprocess
 import tempfile
 import testgres
@@ -13,6 +12,7 @@
 
 from contextlib import contextmanager
 from distutils.version import LooseVersion
+from shutil import rmtree
 
 from testgres import \
     InitNodeException, \
@@ -139,7 +139,7 @@ def test_node_exit(self):
 
         # we should save the DB for "debugging"
         self.assertTrue(os.path.exists(base_dir))
-        shutil.rmtree(base_dir, ignore_errors=True)
+        rmtree(base_dir, ignore_errors=True)
 
         with get_new_node().init() as node:
             base_dir = node.base_dir