postgrespro · demonolock · Dec 19, 2023 · Dec 14, 2023 · Dec 17, 2023 · Dec 17, 2023
diff --git a/testgres/helpers/__init__.py b/testgres/helpers/__init__.py
diff --git a/testgres/helpers/port_manager.py b/testgres/helpers/port_manager.py
@@ -0,0 +1,40 @@
+import socket
+import random
+from typing import Set, Iterable, Optional
+
+
+class PortForException(Exception):
+    pass
+
+
+class PortManager:
+    def __init__(self, ports_range=(1024, 65535)):
+        self.ports_range = ports_range
+
+    @staticmethod
+    def is_port_free(port: int) -> bool:
+        """Check if a port is free to use."""
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            try:
+                s.bind(("", port))
+                return True
+            except OSError:
+                return False
+
+    def find_free_port(self, ports: Optional[Set[int]] = None, exclude_ports: Optional[Iterable[int]] = None) -> int:
+        """Return a random unused port number."""
+        if ports is None:
+            ports = set(range(1024, 65535))
+
+        if exclude_ports is None:
+            exclude_ports = set()
+
+        ports.difference_update(set(exclude_ports))
+
+        sampled_ports = random.sample(tuple(ports), min(len(ports), 100))
+
+        for port in sampled_ports:
+            if self.is_port_free(port):
+                return port
+
+        raise PortForException("Can't select a port")
diff --git a/testgres/operations/local_ops.py b/testgres/operations/local_ops.py
@@ -39,23 +39,30 @@ def __init__(self, conn_params=None):
         self.username = conn_params.username or self.get_user()
 
     @staticmethod
-    def _run_command(cmd, shell, input, timeout, encoding, temp_file=None):
+    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 = temp_file
-            stderr = subprocess.STDOUT
+            stdout = stdout or temp_file
+            stderr = stderr or subprocess.STDOUT
         else:
-            stdout = subprocess.PIPE
-            stderr = subprocess.PIPE
+            stdout = stdout or subprocess.PIPE
+            stderr = stderr or subprocess.PIPE
 
         process = subprocess.Popen(
             cmd,
             shell=shell,
-            stdin=subprocess.PIPE if input is not None else None,
+            stdin=stdin or subprocess.PIPE if input is not None else None,
             stdout=stdout,
             stderr=stderr,
         )
-        return process
+
+        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):
@@ -67,7 +74,7 @@ def _raise_exec_exception(message, command, exit_code, output):
 
     def exec_command(self, cmd, wait_exit=False, verbose=False,
                      expect_error=False, encoding=None, shell=False, text=False,
-                     input=None, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+                     input=None, stdin=None, stdout=None, stderr=None,
                      get_process=None, timeout=None):
         """
         Execute a command in a subprocess.
@@ -143,29 +150,27 @@ def _process_output(process, encoding, temp_file=None):
 
     def _exec_command_windows(self, cmd, wait_exit=False, verbose=False,
                               expect_error=False, encoding=None, shell=False, text=False,
-                              input=None, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+                              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, timeout, encoding, temp_file)
-            try:
-                output = process.communicate(input=input.encode(encoding) if input else None, timeout=timeout)
-            except subprocess.TimeoutExpired:
-                process.kill()
-                raise ExecUtilException("Command timed out after {} seconds.".format(timeout))
+            _, 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(output):
+            if process.returncode != 0 or has_errors(result):
                 if process.returncode == 0:
                     process.returncode = 1
                 if expect_error:
                     if verbose:
-                        return process.returncode, output, output
+                        return process.returncode, result, result
                     else:
-                        return output
+                        return result
                 else:
                     self._raise_exec_exception('Utility exited with non-zero code. Error `{}`', cmd, process.returncode,
-                                               output)
+                                               result)
 
-            return (process.returncode, output, output) if verbose else output
+            return (process.returncode, result, result) if verbose else result
 
     # Environment setup
     def environ(self, var_name):

diff --git a/testgres/utils.py b/testgres/utils.py
@@ -13,7 +13,7 @@
 
 from six import iteritems
 
-from helpers.port_manager import PortManager
+from .helpers.port_manager import PortManager
 from .exceptions import ExecUtilException
 from .config import testgres_config as tconf
 

diff --git a/tests/test_simple.py b/tests/test_simple.py
@@ -141,7 +141,7 @@ def test_init_after_cleanup(self):
             node.cleanup()
             node.init().start().execute('select 1')
 
-    @unittest.skipUnless(util_exists('pg_resetwal'), 'might be missing')
+    @unittest.skipUnless(util_exists('pg_resetwal.exe' if os.name == 'nt' else 'pg_resetwal'), 'pgbench might be missing')
     @unittest.skipUnless(pg_version_ge('9.6'), 'requires 9.6+')
     def test_init_unique_system_id(self):
         # this function exists in PostgreSQL 9.6+
@@ -747,7 +747,7 @@ def test_logging(self):
                 master.restart()
                 self.assertTrue(master._logger.is_alive())
 
-    @unittest.skipUnless(util_exists('pgbench'), 'might be missing')
+    @unittest.skipUnless(util_exists('pgbench.exe' if os.name == 'nt' else 'pgbench'), 'pgbench might be missing')
     def test_pgbench(self):
         with get_new_node().init().start() as node: