PBCKP-588 fix failed tests - psql, set_auto_conf

v.shepard · v.shepard · commit 784b5f818525 · 2023-07-25T17:21:04.000+02:00
diff --git a/testgres/node.py b/testgres/node.py
@@ -3,6 +3,7 @@
 import os
 import random
 import signal
+import subprocess
 import threading
 from queue import Queue
 
@@ -958,19 +959,31 @@ def psql(self,
 
         # select query source
         if query:
-            psql_params.extend(("-c", '"{}"'.format(query)))
+            if self.os_ops.remote:
+                psql_params.extend(("-c", '"{}"'.format(query)))
+            else:
+                psql_params.extend(("-c", query))
         elif filename:
             psql_params.extend(("-f", filename))
         else:
             raise QueryException('Query or filename must be provided')
 
         # should be the last one
         psql_params.append(dbname)
+        if not self.os_ops.remote:
+            # start psql process
+            process = subprocess.Popen(psql_params,
+                                       stdin=subprocess.PIPE,
+                                       stdout=subprocess.PIPE,
+                                       stderr=subprocess.PIPE)
+
+            # wait until it finishes and get stdout and stderr
+            out, err = process.communicate(input=input)
+            return process.returncode, out, err
+        else:
+            status_code, out, err = self.os_ops.exec_command(psql_params, verbose=True, input=input)
 
-        # start psql process
-        status_code, out, err = self.os_ops.exec_command(psql_params, verbose=True, input=input)
-
-        return status_code, out, err
+            return status_code, out, err
 
     @method_decorator(positional_args_hack(['dbname', 'query']))
     def safe_psql(self, query=None, expect_error=False, **kwargs):
@@ -1002,9 +1015,9 @@ def safe_psql(self, query=None, expect_error=False, **kwargs):
             err = e.message
         if ret:
             if expect_error:
-                out = err or b''
+                out = (err or b'').decode('utf-8')
             else:
-                raise QueryException(err or b'', query)
+                raise QueryException((err or b'').decode('utf-8'), query)
         elif expect_error:
             assert False, "Exception was expected, but query finished successfully: `{}` ".format(query)
 
@@ -1529,18 +1542,18 @@ def set_auto_conf(self, options, config='postgresql.auto.conf', rm_options={}):
                                          Defaults to an empty set.
         """
         # parse postgresql.auto.conf
-        auto_conf_file = os.path.join(self.data_dir, config)
-        raw_content = self.os_ops.read(auto_conf_file)
+        path = os.path.join(self.data_dir, config)
 
+        lines = self.os_ops.readlines(path)
         current_options = {}
         current_directives = []
-        for line in raw_content.splitlines():
+        for line in lines:
 
             # ignore comments
             if line.startswith('#'):
                 continue
 
-            if line == '':
+            if line.strip() == '':
                 continue
 
             if line.startswith('include'):
@@ -1570,7 +1583,7 @@ def set_auto_conf(self, options, config='postgresql.auto.conf', rm_options={}):
         for directive in current_directives:
             auto_conf += directive + "\n"
 
-        self.os_ops.write(auto_conf_file, auto_conf)
+        self.os_ops.write(path, auto_conf, truncate=True)
 
 
 class NodeApp:
diff --git a/testgres/operations/local_ops.py b/testgres/operations/local_ops.py
@@ -1,15 +1,15 @@
 import getpass
 import os
 import shutil
+import stat
 import subprocess
 import tempfile
 from shutil import rmtree
 
 import psutil
 
 from ..exceptions import ExecUtilException
-
-from .os_ops import OsOperations, ConnectionParams
+from .os_ops import ConnectionParams, OsOperations
 from .os_ops import pglib
 
 try:
@@ -18,20 +18,24 @@
     from distutils.spawn import find_executable
 
 CMD_TIMEOUT_SEC = 60
+error_markers = [b'error', b'Permission denied', b'fatal']
 
 
 class LocalOperations(OsOperations):
-    def __init__(self, conn_params: ConnectionParams = ConnectionParams()):
-        super().__init__(conn_params.username)
+    def __init__(self, conn_params=None):
+        if conn_params is None:
+            conn_params = ConnectionParams()
+        super(LocalOperations, self).__init__(conn_params.username)
         self.conn_params = conn_params
         self.host = conn_params.host
         self.ssh_key = None
+        self.remote = False
         self.username = conn_params.username or self.get_user()
 
     # Command execution
     def exec_command(self, cmd, wait_exit=False, verbose=False,
-                     expect_error=False, encoding=None, shell=True, text=False,
-                     input=None, stdout=subprocess.PIPE, stderr=subprocess.PIPE, proc=None):
+                     expect_error=False, encoding=None, shell=False, text=False,
+                     input=None, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, proc=None):
         """
         Execute a command in a subprocess.
 
@@ -49,9 +53,6 @@ def exec_command(self, cmd, wait_exit=False, verbose=False,
         - proc: The process to use for subprocess creation.
         :return: The output of the subprocess.
         """
-        if isinstance(cmd, list):
-            cmd = ' '.join(item.decode('utf-8') if isinstance(item, bytes) else item for item in cmd)
-
         if os.name == 'nt':
             with tempfile.NamedTemporaryFile() as buf:
                 process = subprocess.Popen(cmd, stdout=buf, stderr=subprocess.STDOUT)
@@ -71,7 +72,7 @@ def exec_command(self, cmd, wait_exit=False, verbose=False,
             result, error = process.communicate(input)
             exit_status = process.returncode
 
-            found_error = "error" in error.decode(encoding or 'utf-8').lower()
+            error_found = exit_status != 0 or any(marker in error for marker in error_markers)
 
             if encoding:
                 result = result.decode(encoding)
@@ -80,7 +81,7 @@ def exec_command(self, cmd, wait_exit=False, verbose=False,
             if expect_error:
                 raise Exception(result, error)
 
-            if exit_status != 0 or found_error:
+            if exit_status != 0 or error_found:
                 if exit_status == 0:
                     exit_status = 1
                 raise ExecUtilException(message='Utility exited with non-zero code. Error `{}`'.format(error),
@@ -101,7 +102,7 @@ def find_executable(self, executable):
 
     def is_executable(self, file):
         # Check if the file is executable
-        return os.access(file, os.X_OK)
+        return os.stat(file).st_mode & stat.S_IXUSR
 
     def set_env(self, var_name, var_val):
         # Check if the directory is already in PATH
@@ -116,9 +117,12 @@ def get_name(self):
 
     # Work with dirs
     def makedirs(self, path, remove_existing=False):
-        if remove_existing and os.path.exists(path):
-            shutil.rmtree(path)
-        os.makedirs(path, exist_ok=True)
+        if remove_existing:
+            shutil.rmtree(path, ignore_errors=True)
+        try:
+            os.makedirs(path)
+        except FileExistsError:
+            pass
 
     def rmdirs(self, path, ignore_errors=True):
         return rmtree(path, ignore_errors=ignore_errors)
@@ -141,7 +145,7 @@ def pathsep(self):
         return pathsep
 
     def mkdtemp(self, prefix=None):
-        return tempfile.mkdtemp(prefix=prefix)
+        return tempfile.mkdtemp(prefix='{}'.format(prefix))
 
     def mkstemp(self, prefix=None):
         fd, filename = tempfile.mkstemp(prefix=prefix)
diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py
@@ -17,7 +17,7 @@
 sshtunnel.TUNNEL_TIMEOUT = 5.0
 
 
-error_markers = [b'error', b'Permission denied']
+error_markers = [b'error', b'Permission denied', b'fatal']
 
 
 class PsUtilProcessProxy:
@@ -43,6 +43,7 @@ def __init__(self, conn_params: ConnectionParams):
         self.host = conn_params.host
         self.ssh_key = conn_params.ssh_key
         self.ssh = self.ssh_connect()
+        self.remote = True
         self.username = conn_params.username or self.get_user()
         self.tunnel = None
 
@@ -89,7 +90,7 @@ def _read_ssh_key(self):
             ExecUtilException(message="An error occurred while reading the ssh key: {}".format(e))
 
     def exec_command(self, cmd: str, wait_exit=False, verbose=False, expect_error=False,
-                     encoding=None, shell=True, text=False, input=None, stdout=None,
+                     encoding=None, shell=True, text=False, input=None, stdin=None, stdout=None,
                      stderr=None, proc=None):
         """
         Execute a command in the SSH session.
@@ -131,7 +132,11 @@ def exec_command(self, cmd: str, wait_exit=False, verbose=False, expect_error=Fa
         if error_found:
             if exit_status == 0:
                 exit_status = 1
-            raise ExecUtilException(message="Utility exited with non-zero code. Error: {}".format(error.decode(encoding or 'utf-8')),
+            if encoding:
+                message = "Utility exited with non-zero code. Error: {}".format(error.decode(encoding))
+            else:
+                message = b"Utility exited with non-zero code. Error: " + error
+            raise ExecUtilException(message=message,
                                     command=cmd,
                                     exit_code=exit_status,
                                     out=result)
@@ -429,7 +434,7 @@ def db_connect(self, dbname, user, password=None, host="127.0.0.1", port=5432, s
             conn = pglib.connect(
                 host=host,  # change to 'localhost' because we're connecting through a local ssh tunnel
                 port=self.tunnel.local_bind_port,  # use the local bind port set up by the tunnel
-                dbname=dbname,
+                database=dbname,
                 user=user or self.username,
                 password=password
             )
diff --git a/tests/test_remote.py b/tests/test_remote.py
@@ -1,3 +1,5 @@
+import os
+
 import pytest
 
 from testgres import ExecUtilException
@@ -9,9 +11,10 @@ class TestRemoteOperations:
 
     @pytest.fixture(scope="function", autouse=True)
     def setup(self):
-        conn_params = ConnectionParams(host="172.18.0.3",
-                                       username="dev",
-                                       ssh_key='../../container_files/postgres/ssh/id_ed25519')
+        conn_params = ConnectionParams(host=os.getenv('RDBMS_TESTPOOL1_HOST') or '172.18.0.3',
+                                       username='dev',
+                                       ssh_key=os.getenv(
+                                           'RDBMS_TESTPOOL_SSHKEY') or '../../container_files/postgres/ssh/id_ed25519')
         self.operations = RemoteOperations(conn_params)
 
         yield
@@ -35,7 +38,7 @@ def test_exec_command_failure(self):
             exit_status, result, error = self.operations.exec_command(cmd, verbose=True, wait_exit=True)
         except ExecUtilException as e:
             error = e.message
-        assert error == 'Utility exited with non-zero code. Error: bash: line 1: nonexistent_command: command not found\n'
+        assert error == b'Utility exited with non-zero code. Error: bash: line 1: nonexistent_command: command not found\n'
 
     def test_is_executable_true(self):
         """
@@ -62,7 +65,7 @@ def test_makedirs_and_rmdirs_success(self):
         cmd = "pwd"
         pwd = self.operations.exec_command(cmd, wait_exit=True, encoding='utf-8').strip()
 
-        path = f"{pwd}/test_dir"
+        path = "{}/test_dir".format(pwd)
 
         # Test makedirs
         self.operations.makedirs(path)
@@ -88,7 +91,7 @@ def test_makedirs_and_rmdirs_failure(self):
             exit_status, result, error = self.operations.rmdirs(path, verbose=True)
         except ExecUtilException as e:
             error = e.message
-        assert error == "Utility exited with non-zero code. Error: rm: cannot remove '/root/test_dir': Permission denied\n"
+        assert error == b"Utility exited with non-zero code. Error: rm: cannot remove '/root/test_dir': Permission denied\n"
 
     def test_listdir(self):
         """
diff --git a/tests/test_simple_remote.py b/tests/test_simple_remote.py
