Override self.username in both operations

asavchkov · asavchkov · commit e6fc5f132762 · 2024-06-15T23:28:00.000+07:00
diff --git a/testgres/node.py b/testgres/node.py
@@ -63,7 +63,6 @@
 
 from .defaults import \
     default_dbname, \
-    default_username, \
     generate_app_name
 
 from .exceptions import \
@@ -158,7 +157,6 @@ def __init__(self, name=None, port=None, base_dir=None, conn_params: ConnectionP
 
         self.host = self.os_ops.host
         self.port = port or reserve_port()
-        self.username = self.os_ops.username
         self.ssh_key = self.os_ops.ssh_key
 
         # defaults for __exit__()
@@ -692,7 +690,7 @@ def slow_start(self, replica=False, dbname='template1', username=None, max_attem
         # Call poll_query_until until the expected value is returned
         self.poll_query_until(query=query,
                               dbname=dbname,
-                              username=username or self.username,
+                              username=username or self.os_ops.username,
                               suppress={InternalError,
                                         QueryException,
                                         ProgrammingError,
@@ -971,7 +969,7 @@ def psql(self,
             self._get_bin_path("psql"),
             "-p", str(self.port),
             "-h", self.host,
-            "-U", username or self.username,
+            "-U", username or self.os_ops.username,
             "-X",  # no .psqlrc
             "-A",  # unaligned output
             "-t",  # print rows only
@@ -1090,7 +1088,7 @@ def tmpfile():
             "-p", str(self.port),
             "-h", self.host,
             "-f", filename,
-            "-U", username or self.username,
+            "-U", username or self.os_ops.username,
             "-d", dbname or default_dbname(),
             "-F", format.value
         ]  # yapf: disable
@@ -1111,7 +1109,7 @@ def restore(self, filename, dbname=None, username=None):
 
         # Set default arguments
         dbname = dbname or default_dbname()
-        username = username or self.username
+        username = username or self.os_ops.username
 
         _params = [
             self._get_bin_path("pg_restore"),
@@ -1387,7 +1385,7 @@ def pgbench(self,
             self._get_bin_path("pgbench"),
             "-p", str(self.port),
             "-h", self.host,
-            "-U", username or self.username
+            "-U", username or self.os_ops.username
         ] + options  # yapf: disable
 
         # should be the last one
@@ -1460,7 +1458,7 @@ def pgbench_run(self, dbname=None, username=None, options=[], **kwargs):
             self._get_bin_path("pgbench"),
             "-p", str(self.port),
             "-h", self.host,
-            "-U", username or self.username
+            "-U", username or self.os_ops.username
         ] + options  # yapf: disable
 
         for key, value in iteritems(kwargs):
diff --git a/testgres/operations/local_ops.py b/testgres/operations/local_ops.py
@@ -38,6 +38,7 @@ def __init__(self, conn_params=None):
         self.host = conn_params.host
         self.ssh_key = None
         self.remote = False
+        self.username = conn_params.username or getpass.getuser()
 
     @staticmethod
     def _raise_exec_exception(message, command, exit_code, output):
diff --git a/testgres/operations/os_ops.py b/testgres/operations/os_ops.py
@@ -24,7 +24,7 @@ def get_default_encoding():
 class OsOperations:
     def __init__(self, username=None):
         self.ssh_key = None
-        self.username = username or self.get_user()
+        self.username = username or getpass.getuser()
 
     # Command execution
     def exec_command(self, cmd, **kwargs):
@@ -46,7 +46,7 @@ def set_env(self, var_name, var_val):
         raise NotImplementedError()
 
     def get_user(self):
-        return getpass.getuser()
+        return self.username
 
     def get_name(self):
         raise NotImplementedError()
diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py
@@ -1,3 +1,4 @@
+import getpass
 import os
 import socket
 import subprocess
@@ -52,6 +53,7 @@ def __init__(self, conn_params: ConnectionParams):
         if self.port:
             self.ssh_args += ["-p", self.port]
         self.remote = True
+        self.username = conn_params.username or getpass.getuser()
         self.ssh_dest = f"{self.username}@{self.host}" if conn_params.username else self.host
         self.add_known_host(self.host)
         self.tunnel_process = None
