PBCKP-588 test fix test_restore_after_failover

v.shepard · v.shepard · commit 2c2d2c5cf0ea · 2023-06-23T00:36:58.000+02:00
diff --git a/testgres/node.py b/testgres/node.py
@@ -625,9 +625,11 @@ def status(self):
                 "-D", self.data_dir,
                 "status"
             ]  # yapf: disable
-            out = execute_utility(_params, self.utils_log_file)
-            if 'no server running' in out:
+            status_code, out, err = execute_utility(_params, self.utils_log_file, verbose=True)
+            if 'does not exist' in err:
                 return NodeStatus.Uninitialized
+            elif'no server running' in out:
+                return NodeStatus.Stopped
             return NodeStatus.Running
 
         except ExecUtilException as e:
@@ -712,14 +714,17 @@ def start(self, params=[], wait=True):
         ] + params  # yapf: disable
 
         try:
-            execute_utility(_params, self.utils_log_file)
+            exit_status, out, error = execute_utility(_params, self.utils_log_file, verbose=True)
+            if 'does not exist' in error:
+                raise Exception
+            if 'server started' in out:
+                self.is_started = True
         except Exception as e:
             msg = 'Cannot start node'
             files = self._collect_special_files()
             raise_from(StartNodeException(msg, files), e)
 
         self._maybe_start_logger()
-        self.is_started = True
         return self
 
     def stop(self, params=[], wait=True):
diff --git a/testgres/operations/local_ops.py b/testgres/operations/local_ops.py
@@ -52,7 +52,7 @@ def exec_command(self, cmd, wait_exit=False, verbose=False,
         :return: The output of the subprocess.
         """
         if isinstance(cmd, list):
-            cmd = " ".join(cmd)
+            cmd = ' '.join(item.decode('utf-8') if isinstance(item, bytes) else item for item in cmd)
         log.debug(f"Executing command: `{cmd}`")
 
         if os.name == 'nt':
@@ -98,8 +98,7 @@ def exec_command(self, cmd, wait_exit=False, verbose=False,
 
     # Environment setup
     def environ(self, var_name):
-        cmd = f"echo ${var_name}"
-        return self.exec_command(cmd, encoding='utf-8').strip()
+        return os.environ.get(var_name)
 
     def find_executable(self, executable):
         return find_executable(executable)
@@ -108,17 +107,6 @@ def is_executable(self, file):
         # Check if the file is executable
         return os.access(file, os.X_OK)
 
-    def add_to_path(self, new_path):
-        pathsep = self.pathsep
-        # Check if the directory is already in PATH
-        path = self.environ("PATH")
-        if new_path not in path.split(pathsep):
-            if self.remote:
-                self.exec_command(f"export PATH={new_path}{pathsep}{path}")
-            else:
-                os.environ["PATH"] = f"{new_path}{pathsep}{path}"
-        return pathsep
-
     def set_env(self, var_name, var_val):
         # Check if the directory is already in PATH
         os.environ[var_name] = var_val
@@ -128,8 +116,7 @@ def get_user(self):
         return getpass.getuser()
 
     def get_name(self):
-        cmd = 'python3 -c "import os; print(os.name)"'
-        return self.exec_command(cmd).strip()
+        return os.name
 
     # Work with dirs
     def makedirs(self, path, remove_existing=False):
diff --git a/testgres/operations/os_ops.py b/testgres/operations/os_ops.py
@@ -29,9 +29,6 @@ def is_executable(self, file):
         # Check if the file is executable
         raise NotImplementedError()
 
-    def add_to_path(self, new_path):
-        raise NotImplementedError()
-
     def set_env(self, var_name, var_val):
         # Check if the directory is already in PATH
         raise NotImplementedError()
diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py
@@ -20,6 +20,22 @@
 error_markers = [b'error', b'Permission denied']
 
 
+class PsUtilProcessProxy:
+    def __init__(self, ssh, pid):
+        self.ssh = ssh
+        self.pid = pid
+
+    def kill(self):
+        command = f"kill {self.pid}"
+        self.ssh.exec_command(command)
+
+    def cmdline(self):
+        command = f"ps -p {self.pid} -o cmd --no-headers"
+        stdin, stdout, stderr = self.ssh.exec_command(command)
+        cmdline = stdout.read().decode('utf-8').strip()
+        return cmdline.split()
+
+
 class RemoteOperations(OsOperations):
     def __init__(self, host="127.0.0.1", hostname='localhost', port=None, ssh_key=None, username=None):
         super().__init__(username)
@@ -71,7 +87,7 @@ def exec_command(self, cmd: str, wait_exit=False, verbose=False, expect_error=Fa
             self.ssh = self.ssh_connect()
 
         if isinstance(cmd, list):
-            cmd = " ".join(cmd)
+            cmd = ' '.join(item.decode('utf-8') if isinstance(item, bytes) else item for item in cmd)
         if input:
             stdin, stdout, stderr = self.ssh.exec_command(cmd)
             stdin.write(input)
@@ -140,17 +156,6 @@ def is_executable(self, file):
         is_exec = self.exec_command(f"test -x {file} && echo OK")
         return is_exec == b"OK\n"
 
-    def add_to_path(self, new_path):
-        pathsep = self.pathsep
-        # Check if the directory is already in PATH
-        path = self.environ("PATH")
-        if new_path not in path.split(pathsep):
-            if self.remote:
-                self.exec_command(f"export PATH={new_path}{pathsep}{path}")
-            else:
-                os.environ["PATH"] = f"{new_path}{pathsep}{path}"
-        return pathsep
-
     def set_env(self, var_name: str, var_val: str):
         """
         Set the value of an environment variable.
@@ -243,9 +248,17 @@ def mkdtemp(self, prefix=None):
             raise ExecUtilException("Could not create temporary directory.")
 
     def mkstemp(self, prefix=None):
-        cmd = f"mktemp {prefix}XXXXXX"
-        filename = self.exec_command(cmd).strip()
-        return filename
+        if prefix:
+            temp_dir = self.exec_command(f"mktemp {prefix}XXXXX", encoding='utf-8')
+        else:
+            temp_dir = self.exec_command("mktemp", encoding='utf-8')
+
+        if temp_dir:
+            if not os.path.isabs(temp_dir):
+                temp_dir = os.path.join('/home', self.username, temp_dir.strip())
+            return temp_dir
+        else:
+            raise ExecUtilException("Could not create temporary directory.")
 
     def copytree(self, src, dst):
         if not os.path.isabs(dst):
@@ -291,7 +304,7 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal
                 data = data.encode(encoding)
             if isinstance(data, list):
                 # ensure each line ends with a newline
-                data = [s if s.endswith('\n') else s + '\n' for s in data]
+                data = [(s if isinstance(s, str) else s.decode('utf-8')).rstrip('\n') + '\n' for s in data]
                 tmp_file.writelines(data)
             else:
                 tmp_file.write(data)
@@ -351,8 +364,8 @@ def isfile(self, remote_file):
 
     def isdir(self, dirname):
         cmd = f"if [ -d {dirname} ]; then echo True; else echo False; fi"
-        response = self.exec_command(cmd, encoding='utf-8')
-        return response.strip() == "True"
+        response = self.exec_command(cmd)
+        return response.strip() == b"True"
 
     def remove_file(self, filename):
         cmd = f"rm {filename}"
@@ -366,16 +379,16 @@ def kill(self, pid, signal):
 
     def get_pid(self):
         # Get current process id
-        return self.exec_command("echo $$")
+        return int(self.exec_command("echo $$", encoding='utf-8'))
 
     def get_remote_children(self, pid):
         command = f"pgrep -P {pid}"
         stdin, stdout, stderr = self.ssh.exec_command(command)
         children = stdout.readlines()
-        return [int(child_pid.strip()) for child_pid in children]
+        return [PsUtilProcessProxy(self.ssh, int(child_pid.strip())) for child_pid in children]
 
     # Database control
-    def db_connect(self, dbname, user, password=None, host="127.0.0.1", port=5432):
+    def db_connect(self, dbname, user, password=None, host="127.0.0.1", port=5432, ssh_key=None):
         """
         Connects to a PostgreSQL database on the remote system.
         Args:
@@ -389,19 +402,26 @@ def db_connect(self, dbname, user, password=None, host="127.0.0.1", port=5432):
         This function establishes a connection to a PostgreSQL database on the remote system using the specified
         parameters. It returns a connection object that can be used to interact with the database.
         """
-        with sshtunnel.open_tunnel(
-                (host, 22),  # Remote server IP and SSH port
-                ssh_username=self.username,
-                ssh_pkey=self.ssh_key,
-                remote_bind_address=(host, port),  # PostgreSQL server IP and PostgreSQL port
-                local_bind_address=('localhost', port),  # Local machine IP and available port
-        ):
+        tunnel = sshtunnel.open_tunnel(
+            (host, 22),  # Remote server IP and SSH port
+            ssh_username=user or self.username,
+            ssh_pkey=ssh_key or self.ssh_key,
+            remote_bind_address=(host, port),  # PostgreSQL server IP and PostgreSQL port
+            local_bind_address=('localhost', port)  # Local machine IP and available port
+        )
+
+        tunnel.start()
+
+        try:
             conn = pglib.connect(
-                host=host,
-                port=port,
+                host=host,  # change to 'localhost' because we're connecting through a local ssh tunnel
+                port=tunnel.local_bind_port,  # use the local bind port set up by the tunnel
                 dbname=dbname,
-                user=user,
+                user=user or self.username,
                 password=password
             )
 
-        return conn
+            return conn
+        except Exception as e:
+            tunnel.stop()
+            raise e
diff --git a/testgres/utils.py b/testgres/utils.py
@@ -45,7 +45,7 @@ def release_port(port):
     bound_ports.discard(port)
 
 
-def execute_utility(args, logfile=None):
+def execute_utility(args, logfile=None, verbose=False):
     """
     Execute utility (pg_ctl, pg_dump etc).
 
@@ -56,24 +56,28 @@ def execute_utility(args, logfile=None):
     Returns:
         stdout of executed utility.
     """
-    command = u' '.join(args)
-    exit_status, out, error = tconf.os_ops.exec_command(command, verbose=True)
+    exit_status, out, error = tconf.os_ops.exec_command(args, verbose=True)
     # decode result
     out = '' if not out else out
     if isinstance(out, bytes):
         out = out.decode('utf-8')
+    if isinstance(error, bytes):
+        error = error.decode('utf-8')
 
     # write new log entry if possible
     if logfile:
         try:
-            tconf.os_ops.write(filename=logfile, data=command, truncate=True)
+            tconf.os_ops.write(filename=logfile, data=args, truncate=True)
             if out:
                 # comment-out lines
                 lines = [u'\n'] + ['# ' + line for line in out.splitlines()] + [u'\n']
                 tconf.os_ops.write(filename=logfile, data=lines)
         except IOError:
-            log.warn(f"Problem with writing to logfile `{logfile}` during run command `{command}`")
-    return out
+            log.warn(f"Problem with writing to logfile `{logfile}` during run command `{args}`")
+    if verbose:
+        return exit_status, out, error
+    else:
+        return out
 
 
 def get_bin_path(filename):
