PBCKP-588 fixes after review - add ConnectionParams

v.shepard · v.shepard · commit 0528541e70a3 · 2023-06-26T02:23:41.000+02:00
diff --git a/testgres/__init__.py b/testgres/__init__.py
@@ -46,7 +46,7 @@
     First, \
     Any
 
-from .operations.os_ops import OsOperations
+from .operations.os_ops import OsOperations, ConnectionParams
 from .operations.local_ops import LocalOperations
 from .operations.remote_ops import RemoteOperations
 
@@ -60,5 +60,5 @@
     "PostgresNode", "NodeApp",
     "reserve_port", "release_port", "bound_ports", "get_bin_path", "get_pg_config", "get_pg_version",
     "First", "Any",
-    "OsOperations", "LocalOperations", "RemoteOperations"
+    "OsOperations", "LocalOperations", "RemoteOperations", "ConnectionParams"
 ]
diff --git a/testgres/api.py b/testgres/api.py
@@ -37,11 +37,10 @@ def get_new_node(name=None, base_dir=None, **kwargs):
     """
     Simply a wrapper around :class:`.PostgresNode` constructor.
     See :meth:`.PostgresNode.__init__` for details.
-    For remote connection you can add next parameters:
-    host='127.0.0.1',
-    hostname='localhost',
-    ssh_key=None,
-    username=default_username()
+    For remote connection you can add the next parameter:
+    conn_params = ConnectionParams(host='127.0.0.1',
+                                   ssh_key=None,
+                                   username=default_username())
     """
     # NOTE: leave explicit 'name' and 'base_dir' for compatibility
     return PostgresNode(name=name, base_dir=base_dir, **kwargs)
diff --git a/testgres/backup.py b/testgres/backup.py
@@ -139,7 +139,7 @@ def spawn_primary(self, name=None, destroy=True):
 
         # Build a new PostgresNode
         NodeClass = self.original_node.__class__
-        with clean_on_error(NodeClass(name=name, base_dir=base_dir, os_ops=self.original_node.os_ops)) as node:
+        with clean_on_error(NodeClass(name=name, base_dir=base_dir, conn_params=self.original_node.os_ops.conn_params)) as node:
 
             # New nodes should always remove dir tree
             node._should_rm_dirs = True
diff --git a/testgres/node.py b/testgres/node.py
@@ -94,6 +94,7 @@
 
 from .backup import NodeBackup
 
+from .operations.os_ops import ConnectionParams
 from .operations.local_ops import LocalOperations
 from .operations.remote_ops import RemoteOperations
 
@@ -125,8 +126,7 @@ def __repr__(self):
 
 
 class PostgresNode(object):
-    def __init__(self, name=None, port=None, base_dir=None,
-                 host='127.0.0.1', hostname='localhost', ssh_key=None, username=default_username(), os_ops=None):
+    def __init__(self, name=None, port=None, base_dir=None, conn_params: ConnectionParams = ConnectionParams()):
         """
         PostgresNode constructor.
 
@@ -146,17 +146,14 @@ def __init__(self, name=None, port=None, base_dir=None,
         # basic
         self.name = name or generate_app_name()
 
-        if os_ops:
-            self.os_ops = os_ops
-        elif ssh_key:
-            self.os_ops = RemoteOperations(host=host, hostname=hostname, ssh_key=ssh_key, username=username)
+        if conn_params.ssh_key:
+            self.os_ops = RemoteOperations(conn_params)
         else:
-            self.os_ops = LocalOperations(host=host, hostname=hostname, username=username)
+            self.os_ops = LocalOperations(conn_params)
 
-        self.port = self.os_ops.port or reserve_port()
+        self.port = port or reserve_port()
 
         self.host = self.os_ops.host
-        self.hostname = self.os_ops.hostname
         self.ssh_key = self.os_ops.ssh_key
 
         testgres_config.os_ops = self.os_ops
@@ -628,7 +625,7 @@ def status(self):
             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:
+            elif 'no server running' in out:
                 return NodeStatus.Stopped
             return NodeStatus.Running
 
diff --git a/testgres/operations/local_ops.py b/testgres/operations/local_ops.py
@@ -7,9 +7,9 @@
 
 import psutil
 
-from testgres.exceptions import ExecUtilException
+from ..exceptions import ExecUtilException
 
-from .os_ops import OsOperations
+from .os_ops import OsOperations, ConnectionParams
 from .os_ops import pglib
 
 try:
@@ -21,13 +21,12 @@
 
 
 class LocalOperations(OsOperations):
-    def __init__(self, host='127.0.0.1', hostname='localhost', port=None, username=None):
-        super().__init__(username)
-        self.host = host
-        self.hostname = hostname
-        self.port = port
+    def __init__(self, conn_params: ConnectionParams = ConnectionParams()):
+        super().__init__(conn_params.username)
+        self.conn_params = conn_params
+        self.host = conn_params.host
         self.ssh_key = None
-        self.username = username or self.get_user()
+        self.username = conn_params.username or self.get_user()
 
     # Command execution
     def exec_command(self, cmd, wait_exit=False, verbose=False,
diff --git a/testgres/operations/os_ops.py b/testgres/operations/os_ops.py
@@ -7,11 +7,16 @@
         raise ImportError("You must have psycopg2 or pg8000 modules installed")
 
 
+class ConnectionParams:
+    def __init__(self, host='127.0.0.1', ssh_key=None, username=None):
+        self.host = host
+        self.ssh_key = ssh_key
+        self.username = username
+
+
 class OsOperations:
     def __init__(self, username=None):
-        self.hostname = "localhost"
-        self.remote = False
-        self.ssh = None
+        self.ssh_key = None
         self.username = username
 
     # Command execution
diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py
@@ -8,9 +8,9 @@
 import paramiko
 from paramiko import SSHClient
 
-from testgres.exceptions import ExecUtilException
+from ..exceptions import ExecUtilException
 
-from .os_ops import OsOperations
+from .os_ops import OsOperations, ConnectionParams
 from .os_ops import pglib
 
 sshtunnel.SSH_TIMEOUT = 5.0
@@ -37,15 +37,13 @@ def cmdline(self):
 
 
 class RemoteOperations(OsOperations):
-    def __init__(self, host="127.0.0.1", hostname='localhost', port=None, ssh_key=None, username=None):
-        super().__init__(username)
-        self.host = host
-        self.hostname = hostname
-        self.port = port
-        self.ssh_key = ssh_key
-        self.remote = True
+    def __init__(self, conn_params: ConnectionParams):
+        super().__init__(conn_params.username)
+        self.conn_params = conn_params
+        self.host = conn_params.host
+        self.ssh_key = conn_params.ssh_key
         self.ssh = self.ssh_connect()
-        self.username = username or self.get_user()
+        self.username = conn_params.username or self.get_user()
         self.tunnel = None
 
     def __enter__(self):
@@ -70,14 +68,11 @@ def close_tunnel(self):
                 time.sleep(0.5)
 
     def ssh_connect(self) -> Optional[SSHClient]:
-        if not self.remote:
-            return None
-        else:
-            key = self._read_ssh_key()
-            ssh = paramiko.SSHClient()
-            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-            ssh.connect(self.host, username=self.username, pkey=key)
-            return ssh
+        key = self._read_ssh_key()
+        ssh = paramiko.SSHClient()
+        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        ssh.connect(self.host, username=self.username, pkey=key)
+        return ssh
 
     def _read_ssh_key(self):
         try:
diff --git a/tests/test_remote.py b/tests/test_remote.py
@@ -2,20 +2,19 @@
 
 from testgres import ExecUtilException
 from testgres import RemoteOperations
+from testgres import ConnectionParams
 
 
 class TestRemoteOperations:
 
     @pytest.fixture(scope="function", autouse=True)
     def setup(self):
-        self.operations = RemoteOperations(
-            host="172.18.0.3",
-            username="dev",
-            ssh_key='../../container_files/postgres/ssh/id_ed25519'
-        )
+        conn_params = ConnectionParams(host="172.18.0.3",
+                                       username="dev",
+                                       ssh_key='../../container_files/postgres/ssh/id_ed25519')
+        self.operations = RemoteOperations(conn_params)
 
         yield
-
         self.operations.__del__()
 
     def test_exec_command_success(self):
diff --git a/tests/test_simple_remote.py b/tests/test_simple_remote.py
@@ -50,12 +50,12 @@
 # NOTE: those are ugly imports
 from testgres import bound_ports
 from testgres.utils import PgVer
-from testgres.node import ProcessProxy
+from testgres.node import ProcessProxy, ConnectionParams
 
-
-os_ops = RemoteOperations(host='172.18.0.3',
-                          username='dev',
-                          ssh_key='../../container_files/postgres/ssh/id_ed25519')
+conn_params = ConnectionParams(host="172.18.0.3",
+                               username="dev",
+                               ssh_key='../../container_files/postgres/ssh/id_ed25519')
+os_ops = RemoteOperations(conn_params)
 testgres_config.set_os_ops(os_ops=os_ops)
 
 
@@ -94,7 +94,7 @@ def removing(f):
 
 
 def get_remote_node(name=None):
-    return get_new_node(name=name, host=os_ops.host, username=os_ops.username, ssh_key=os_ops.ssh_key)
+    return get_new_node(name=name, conn_params=conn_params)
 
 
 class TestgresRemoteTests(unittest.TestCase):
