Update establish_ssh_tunnel

vshepard · vshepard · commit 1eb9a92fd7fe · 2024-06-06T08:18:49.000+02:00
diff --git a/setup.py b/setup.py
@@ -27,7 +27,7 @@
     readme = f.read()
 
 setup(
-    version='1.10.1',
+    version='1.10.2',
     name='testgres',
     packages=['testgres', 'testgres.operations', 'testgres.helpers'],
     description='Testing utility for PostgreSQL and its extensions',
diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py
@@ -76,11 +76,14 @@ def is_port_open(host, port):
             except socket.error:
                 return False
 
-    def establish_ssh_tunnel(self, local_port, remote_port):
+    def establish_ssh_tunnel(self, local_port, remote_port, host):
         """
         Establish an SSH tunnel from a local port to a remote PostgreSQL port.
         """
-        ssh_cmd = ['-N', '-L', f"{local_port}:localhost:{remote_port}"]
+        if host != 'localhost':
+            ssh_cmd = ['-N', '-L', f"localhost:{local_port}:{host}:{remote_port}"]
+        else:
+            ssh_cmd = ['-N', '-L', f"{local_port}:{host}:{remote_port}"]
         self.tunnel_process = self.exec_command(ssh_cmd, get_process=True, timeout=300)
         timeout = 10
         start_time = time.time()
@@ -412,10 +415,10 @@ def db_connect(self, dbname, user, password=None, host="localhost", port=5432):
         """
         local_port = reserve_port()
         self.tunnel_port = local_port
-        self.establish_ssh_tunnel(local_port=local_port, remote_port=port)
+        self.establish_ssh_tunnel(local_port=local_port, remote_port=port, host=host)
         try:
             conn = pglib.connect(
-                host=host,
+                host='localhost',
                 port=local_port,
                 database=dbname,
                 user=user,
