Review fixes

cmapi/integration_tests/ssh.py

@@ -1,22 +1,22 @@
 import concurrent.futures
-from contextlib import contextmanager
 import json
+from contextlib import contextmanager
 from pathlib import Path
 from typing import Any, Callable
+
 from fabric import Connection
 
 
-class HostConfig:
-    """Keeps configuration for a host in the cluster"""
-    def __init__(self, name: str, private_ip: str, public_fqdn: str, key_file_path: Path, ssh_user: str) -> None:
-        self.name = name
-        self.private_ip = private_ip
-        self.public_fqdn = public_fqdn
-        self.key_file_path = key_file_path
-        self.ssh_user = ssh_user
+from dataclasses import dataclass, field
 
-    def __repr__(self) -> str:
-        return f"HostConfig(name={self.name}, private_ip={self.private_ip}, public_fqdn={self.public_fqdn}"
+@dataclass
+class RemoteHost:
+    """Keeps configuration for a host in the cluster"""
+    name: str
+    private_ip: str
+    public_fqdn: str
+    key_file_path: Path = field(repr=False)
+    ssh_user: str = field(repr=False)
 
     @contextmanager
     def ssh_connection(self):
@@ -56,7 +56,7 @@ class HostConfig:
 
 class ClusterConfig:
     """Keeps configuration of the cluster"""
-    def __init__(self, mcs_hosts: list[HostConfig], maxscale_hosts: list[HostConfig]) -> None:
+    def __init__(self, mcs_hosts: list[RemoteHost], maxscale_hosts: list[RemoteHost]) -> None:
         self.mcs_hosts = mcs_hosts
         self.maxscale_hosts = maxscale_hosts
 
@@ -64,25 +64,28 @@ class ClusterConfig:
         return f"ClusterConfig(mcs_hosts={self.mcs_hosts}"
 
     @property
-    def primary(self) -> HostConfig:
+    def primary(self) -> RemoteHost:
         return self.mcs_hosts[0]
 
     @property
-    def replicas(self) -> list[HostConfig]:
+    def replicas(self) -> list[RemoteHost]:
         return self.mcs_hosts[1:]
 
 
-def run_on_all_hosts_parallel(hosts: list[HostConfig], func: Callable[[HostConfig], Any], max_workers=None):
+def run_on_all_hosts_parallel(
+      hosts: list[RemoteHost],
+      func: Callable[[RemoteHost], Any],
+      timeout: float | None = None,
+      max_workers: int | None = None,
+    ) -> dict[str, Any]:
     """
     Run a function on all hosts in parallel.
 
-    Args:
-        hosts: List of HostConfig objects
-        func: Function that takes a HostConfig as its only argument
-        max_workers: Maximum number of worker threads (None = default based on system)
-
-    Returns:
-        Dictionary mapping hosts to function results
+    :param hosts: List of RemoteHost objects
+    :param func: Function that takes a RemoteHost as its only argument
+    :param timeout: Timeout for the function execution
+    :param max_workers: Maximum number of worker threads (None = default based on system)
+    :returns: Dictionary mapping hosts to function results
     """
     results = {}
 
@@ -94,7 +97,7 @@ def run_on_all_hosts_parallel(hosts: list[HostConfig], func: Callable[[HostConfi
         for future in concurrent.futures.as_completed(future_to_host):
             host = future_to_host[future]
             try:
-                result = future.result()
+                result = future.result(timeout=timeout)
                 results[host.name] = result
             except Exception as exc:
                 print(f"Error on host {host.name}: {exc}")
@@ -104,7 +107,7 @@ def run_on_all_hosts_parallel(hosts: list[HostConfig], func: Callable[[HostConfi
     return results
 
 
-def block_port(host: HostConfig, port: int) -> None:
+def block_port(host: RemoteHost, port: int) -> None:
     """Block a port on the host using iptables"""
     with host.ssh_connection() as conn:
         print(f"{host.name}: Blocking port {port}")
@@ -112,7 +115,7 @@ def block_port(host: HostConfig, port: int) -> None:
         print(f"{host.name}: Port {port} blocked")
 
 
-def unblock_port(host: HostConfig, port: int) -> None:
+def unblock_port(host: RemoteHost, port: int) -> None:
     """Unblock a port on the host using iptables"""
     with host.ssh_connection() as conn:
         print(f"{host.name}: Unblocking port {port}")