Fix initdb error on Windows - fix pgbench

vshepard · vshepard · commit 05cd996c9864 · 2023-12-17T23:43:02.000+01:00
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
@@ -55,12 +55,7 @@ def _run_command(cmd, shell, input, timeout, encoding, temp_file=None):
             stdout=stdout,
             stderr=stderr,
         )
-
-        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))
+        return process
 
     @staticmethod
     def _raise_exec_exception(message, command, exit_code, output):
@@ -151,10 +146,12 @@ def _exec_command_windows(self, cmd, wait_exit=False, verbose=False,
                               input=None, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                               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)
-            if get_process:
-                return process
-            output = self._process_output(process, encoding, 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))
 
             if process.returncode != 0 or has_errors(output):
                 if process.returncode == 0:
diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py
@@ -45,6 +45,10 @@ def __init__(self, conn_params: ConnectionParams):
         self.conn_params = conn_params
         self.host = conn_params.host
         self.ssh_key = conn_params.ssh_key
+        if self.ssh_key:
+            self.ssh_cmd = ["-i", self.ssh_key]
+        else:
+            self.ssh_cmd = []
         self.remote = True
         self.username = conn_params.username or self.get_user()
         self.add_known_host(self.host)
@@ -91,9 +95,9 @@ def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False,
         """
         ssh_cmd = []
         if isinstance(cmd, str):
-            ssh_cmd = ['ssh', f"{self.username}@{self.host}", '-i', self.ssh_key, cmd]
+            ssh_cmd = ['ssh', f"{self.username}@{self.host}"] + self.ssh_cmd + [cmd]
         elif isinstance(cmd, list):
-            ssh_cmd = ['ssh', f"{self.username}@{self.host}", '-i', self.ssh_key] + cmd
+            ssh_cmd = ['ssh', f"{self.username}@{self.host}"] + self.ssh_cmd + cmd
         process = subprocess.Popen(ssh_cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
         if get_process:
             return process
@@ -242,9 +246,9 @@ def mkdtemp(self, prefix=None):
         - prefix (str): The prefix of the temporary directory name.
         """
         if prefix:
-            command = ["ssh", "-i", self.ssh_key, f"{self.username}@{self.host}", f"mktemp -d {prefix}XXXXX"]
+            command = ["ssh"] + self.ssh_cmd + [f"{self.username}@{self.host}", f"mktemp -d {prefix}XXXXX"]
         else:
-            command = ["ssh", "-i", self.ssh_key, f"{self.username}@{self.host}", "mktemp -d"]
+            command = ["ssh"] + self.ssh_cmd + [f"{self.username}@{self.host}", "mktemp -d"]
 
         result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
 
@@ -288,7 +292,7 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal
 
         with tempfile.NamedTemporaryFile(mode=mode, delete=False) as tmp_file:
             if not truncate:
-                scp_cmd = ['scp', '-i', self.ssh_key, f"{self.username}@{self.host}:{filename}", tmp_file.name]
+                scp_cmd = ['scp'] + self.ssh_cmd + [f"{self.username}@{self.host}:{filename}", tmp_file.name]
                 subprocess.run(scp_cmd, check=False)  # The file might not exist yet
                 tmp_file.seek(0, os.SEEK_END)
 
@@ -304,12 +308,11 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal
                 tmp_file.write(data)
 
             tmp_file.flush()
-
-            scp_cmd = ['scp', '-i', self.ssh_key, tmp_file.name, f"{self.username}@{self.host}:{filename}"]
+            scp_cmd = ['scp'] + self.ssh_cmd + [tmp_file.name, f"{self.username}@{self.host}:{filename}"]
             subprocess.run(scp_cmd, check=True)
 
             remote_directory = os.path.dirname(filename)
-            mkdir_cmd = ['ssh', '-i', self.ssh_key, f"{self.username}@{self.host}", f"mkdir -p {remote_directory}"]
+            mkdir_cmd = ['ssh'] + self.ssh_cmd + [f"{self.username}@{self.host}", f"mkdir -p {remote_directory}"]
             subprocess.run(mkdir_cmd, check=True)
 
             os.remove(tmp_file.name)
@@ -374,7 +377,7 @@ def get_pid(self):
         return int(self.exec_command("echo $$", encoding=get_default_encoding()))
 
     def get_process_children(self, pid):
-        command = ["ssh", "-i", self.ssh_key, f"{self.username}@{self.host}", f"pgrep -P {pid}"]
+        command = ["ssh"] + self.ssh_cmd + [f"{self.username}@{self.host}", f"pgrep -P {pid}"]
 
         result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
 
diff --git a/testgres/utils.py b/testgres/utils.py
@@ -4,18 +4,16 @@
 from __future__ import print_function
 
 import os
-import random
-import socket
 
 import sys
 
 from contextlib import contextmanager
 from packaging.version import Version, InvalidVersion
 import re
 
-from port_for import PortForException
 from six import iteritems
 
+from helpers.port_manager import PortManager
 from .exceptions import ExecUtilException
 from .config import testgres_config as tconf
 
@@ -40,49 +38,13 @@ def reserve_port():
     """
     Generate a new port and add it to 'bound_ports'.
     """
-    port = select_random(exclude_ports=bound_ports)
+    port_mng = PortManager()
+    port = port_mng.find_free_port(exclude_ports=bound_ports)
     bound_ports.add(port)
 
     return port
 
 
-def select_random(
-        ports=None,
-        exclude_ports=None,
-) -> int:
-    """
-    Return random unused port number.
-    Standard function from port_for does not work on Windows because of error
-    'port_for.exceptions.PortForException: Can't select a port'
-    We should update it.
-    """
-    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 is_port_free(port):
-            return port
-
-    raise PortForException("Can't select a port")
-
-
-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 release_port(port):
     """
     Free port provided by reserve_port().
diff --git a/tests/test_remote.py b/tests/test_remote.py
@@ -11,10 +11,9 @@ class TestRemoteOperations:
 
     @pytest.fixture(scope="function", autouse=True)
     def setup(self):
-        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')
+        conn_params = ConnectionParams(host=os.getenv('RDBMS_TESTPOOL1_HOST') or '127.0.0.1',
+                                       username=os.getenv('USER'),
+                                       ssh_key=os.getenv('RDBMS_TESTPOOL_SSHKEY'))
         self.operations = RemoteOperations(conn_params)
 
     def test_exec_command_success(self):
@@ -41,7 +40,7 @@ def test_is_executable_true(self):
         """
         Test is_executable for an existing executable.
         """
-        cmd = "postgres"
+        cmd = os.getenv('PG_CONFIG')
         response = self.operations.is_executable(cmd)
 
         assert response is True
diff --git a/tests/test_simple.py b/tests/test_simple.py
@@ -763,6 +763,8 @@ def test_pgbench(self):
             out, _ = proc.communicate()
             out = out.decode('utf-8')
 
+            proc.stdout.close()
+
             self.assertTrue('tps' in out)
 
     def test_pg_config(self):
diff --git a/tests/test_simple_remote.py b/tests/test_simple_remote.py
@@ -52,10 +52,9 @@
 from testgres.utils import PgVer
 from testgres.node import ProcessProxy, ConnectionParams
 
-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')
+conn_params = ConnectionParams(host=os.getenv('RDBMS_TESTPOOL1_HOST') or '127.0.0.1',
+                               username=os.getenv('USER'),
+                               ssh_key=os.getenv('RDBMS_TESTPOOL_SSHKEY'))
 os_ops = RemoteOperations(conn_params)
 testgres_config.set_os_ops(os_ops=os_ops)
 
