Refactoring local_ops.py

demonolock · vshepard · commit 36485874d8e1 · 2023-12-18T11:01:02.000+01:00
diff --git a/setup.py b/setup.py
@@ -29,7 +29,7 @@
 setup(
     version='1.9.3',
     name='testgres',
-    packages=['testgres', 'testgres.operations'],
+    packages=['testgres', 'testgres.operations', 'testgres.helpers'],
     description='Testing utility for PostgreSQL and its extensions',
     url='https://github.com/postgrespro/testgres',
     long_description=readme,
diff --git a/testgres/__init__.py b/testgres/__init__.py
@@ -52,6 +52,8 @@
 from .operations.local_ops import LocalOperations
 from .operations.remote_ops import RemoteOperations
 
+from .helpers.port_manager import PortManager
+
 __all__ = [
     "get_new_node",
     "get_remote_node",
@@ -62,6 +64,6 @@
     "XLogMethod", "IsolationLevel", "NodeStatus", "ProcessType", "DumpFormat",
     "PostgresNode", "NodeApp",
     "reserve_port", "release_port", "bound_ports", "get_bin_path", "get_pg_config", "get_pg_version",
-    "First", "Any",
+    "First", "Any", "PortManager",
     "OsOperations", "LocalOperations", "RemoteOperations", "ConnectionParams"
 ]
diff --git a/testgres/node.py b/testgres/node.py
@@ -623,8 +623,8 @@ def status(self):
                 "-D", self.data_dir,
                 "status"
             ]  # yapf: disable
-            status_code, out, err = execute_utility(_params, self.utils_log_file, verbose=True)
-            if 'does not exist' in err:
+            status_code, out, error = execute_utility(_params, self.utils_log_file, verbose=True)
+            if error and 'does not exist' in error:
                 return NodeStatus.Uninitialized
             elif 'no server running' in out:
                 return NodeStatus.Stopped
@@ -717,7 +717,7 @@ def start(self, params=[], wait=True):
 
         try:
             exit_status, out, error = execute_utility(_params, self.utils_log_file, verbose=True)
-            if 'does not exist' in error:
+            if error and 'does not exist' in error:
                 raise Exception
         except Exception as e:
             msg = 'Cannot start node'
@@ -791,7 +791,7 @@ def restart(self, params=[]):
 
         try:
             error_code, out, error = execute_utility(_params, self.utils_log_file, verbose=True)
-            if 'could not start server' in error:
+            if error and 'could not start server' in error:
                 raise ExecUtilException
         except ExecUtilException as e:
             msg = 'Cannot restart node'
diff --git a/testgres/operations/local_ops.py b/testgres/operations/local_ops.py
@@ -38,32 +38,6 @@ def __init__(self, conn_params=None):
         self.remote = False
         self.username = conn_params.username or self.get_user()
 
-    @staticmethod
-    def _run_command(cmd, shell, input, stdin, stdout, stderr, timeout, encoding, temp_file=None, get_process=None):
-        """Execute a command and return the process."""
-        if temp_file is not None:
-            stdout = stdout or temp_file
-            stderr = stderr or subprocess.STDOUT
-        else:
-            stdout = stdout or subprocess.PIPE
-            stderr = stderr or subprocess.PIPE
-
-        process = subprocess.Popen(
-            cmd,
-            shell=shell,
-            stdin=stdin or subprocess.PIPE if input is not None else None,
-            stdout=stdout,
-            stderr=stderr,
-        )
-
-        if get_process:
-            return None, process
-        try:
-            return process.communicate(input=input.encode(encoding) if input else None, timeout=timeout), process
-        except subprocess.TimeoutExpired:
-            process.kill()
-            raise ExecUtilException("Command timed out after {} seconds.".format(timeout))
-
     @staticmethod
     def _raise_exec_exception(message, command, exit_code, output):
         """Raise an ExecUtilException."""
@@ -72,105 +46,72 @@ def _raise_exec_exception(message, command, exit_code, output):
                                 exit_code=exit_code,
                                 out=output)
 
-    def exec_command(self, cmd, wait_exit=False, verbose=False,
-                     expect_error=False, encoding=None, shell=False, text=False,
-                     input=None, stdin=None, stdout=None, stderr=None,
-                     get_process=None, timeout=None):
-        """
-        Execute a command in a subprocess.
-
-        Args:
-        - cmd: The command to execute.
-        - wait_exit: Whether to wait for the subprocess to exit before returning.
-        - verbose: Whether to return verbose output.
-        - expect_error: Whether to raise an error if the subprocess exits with an error status.
-        - encoding: The encoding to use for decoding the subprocess output.
-        - shell: Whether to use shell when executing the subprocess.
-        - text: Whether to return str instead of bytes for the subprocess output.
-        - input: The input to pass to the subprocess.
-        - stdout: The stdout to use for the subprocess.
-        - stderr: The stderr to use for the subprocess.
-        - proc: The process to use for subprocess creation.
-        :return: The output of the subprocess.
-        """
-        if os.name == 'nt':
-            return self._exec_command_windows(cmd, wait_exit=wait_exit, verbose=verbose,
-                                              expect_error=expect_error, encoding=encoding, shell=shell, text=text,
-                                              input=input, stdin=stdin, stdout=stdout, stderr=stderr,
-                                              get_process=get_process, timeout=timeout)
-        else:
+    @staticmethod
+    def _process_output(encoding, temp_file_path):
+        """Process the output of a command from a temporary file."""
+        with open(temp_file_path, 'rb') as temp_file:
+            output = temp_file.read()
+            if encoding:
+                output = output.decode(encoding)
+            return output, None  # In Windows stderr writing in stdout
+
+    def _run_command(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding):
+        """Execute a command and return the process and its output."""
+        if os.name == 'nt' and stdout is None:  # Windows
+            with tempfile.NamedTemporaryFile(mode='w+b', delete=False) as temp_file:
+                stdout = temp_file
+                stderr = subprocess.STDOUT
+                process = subprocess.Popen(
+                    cmd,
+                    shell=shell,
+                    stdin=stdin or subprocess.PIPE if input is not None else None,
+                    stdout=stdout,
+                    stderr=stderr,
+                )
+                if get_process:
+                    return process, None, None
+                temp_file_path = temp_file.name
+
+            # Wait process finished
+            process.wait()
+
+            output, error = self._process_output(encoding, temp_file_path)
+            return process, output, error
+        else:  # Other OS
             process = subprocess.Popen(
                 cmd,
                 shell=shell,
-                stdin=stdin,
-                stdout=stdout,
-                stderr=stderr,
+                stdin=stdin or subprocess.PIPE if input is not None else None,
+                stdout=stdout or subprocess.PIPE,
+                stderr=stderr or subprocess.PIPE,
             )
             if get_process:
-                return process
-
+                return process, None, None
             try:
-                result, error = process.communicate(input, timeout=timeout)
+                output, error = process.communicate(input=input.encode(encoding) if input else None, timeout=timeout)
+                if encoding:
+                    output = output.decode(encoding)
+                    error = error.decode(encoding)
+                return process, output, error
             except subprocess.TimeoutExpired:
                 process.kill()
                 raise ExecUtilException("Command timed out after {} seconds.".format(timeout))
-            exit_status = process.returncode
 
-            error_found = exit_status != 0 or has_errors(error)
-
-            if encoding:
-                result = result.decode(encoding)
-                error = error.decode(encoding)
-
-            if expect_error:
-                raise Exception(result, error)
-
-            if exit_status != 0 or error_found:
-                if exit_status == 0:
-                    exit_status = 1
-                self._raise_exec_exception('Utility exited with non-zero code. Error `{}`', cmd, exit_status, result)
-            if verbose:
-                return exit_status, result, error
-            else:
-                return result
+    def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False, encoding=None, shell=False,
+                     text=False, input=None, stdin=None, stdout=None, stderr=None, get_process=False, timeout=None):
+        """
+        Execute a command in a subprocess and handle the output based on the provided parameters.
+        """
+        process, output, error = self._run_command(cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding)
+        if get_process:
+            return process
+        if process.returncode != 0 or (has_errors(error) and not expect_error):
+            self._raise_exec_exception('Utility exited with non-zero code. Error `{}`', cmd, process.returncode, error)
 
-    @staticmethod
-    def _process_output(process, encoding, temp_file=None):
-        """Process the output of a command."""
-        if temp_file is not None:
-            temp_file.seek(0)
-            output = temp_file.read()
+        if verbose:
+            return process.returncode, output, error
         else:
-            output = process.stdout.read()
-
-        if encoding:
-            output = output.decode(encoding)
-
-        return output
-
-    def _exec_command_windows(self, cmd, wait_exit=False, verbose=False,
-                              expect_error=False, encoding=None, shell=False, text=False,
-                              input=None, stdin=None, stdout=None, stderr=None,
-                              get_process=None, timeout=None):
-        with tempfile.NamedTemporaryFile(mode='w+b') as temp_file:
-            _, process = self._run_command(cmd, shell, input, stdin, stdout, stderr, timeout, encoding, temp_file, get_process)
-            if get_process:
-                return process
-            result = self._process_output(process, encoding, temp_file)
-
-            if process.returncode != 0 or has_errors(result):
-                if process.returncode == 0:
-                    process.returncode = 1
-                if expect_error:
-                    if verbose:
-                        return process.returncode, result, result
-                    else:
-                        return result
-                else:
-                    self._raise_exec_exception('Utility exited with non-zero code. Error `{}`', cmd, process.returncode,
-                                               result)
-
-            return (process.returncode, result, result) if verbose else result
+            return output
 
     # Environment setup
     def environ(self, var_name):
diff --git a/testgres/operations/os_ops.py b/testgres/operations/os_ops.py
@@ -81,7 +81,7 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal
     def touch(self, filename):
         raise NotImplementedError()
 
-    def read(self, filename):
+    def read(self, filename, encoding, binary):
         raise NotImplementedError()
 
     def readlines(self, filename):
