Implement get_user

asavchkov · asavchkov · commit bc9a62e23258 · 2024-06-15T21:57:21.000+07:00
diff --git a/testgres/node.py b/testgres/node.py
@@ -156,9 +156,10 @@ def __init__(self, name=None, port=None, base_dir=None, conn_params: ConnectionP
         else:
             self.os_ops = LocalOperations(conn_params)
 
-        self.port = port or reserve_port()
-
         self.host = self.os_ops.host
+        self.port = port or reserve_port()
+        # Default node username
+        self.username = default_username()
         self.ssh_key = self.os_ops.ssh_key
 
         # defaults for __exit__()
@@ -683,8 +684,6 @@ def slow_start(self, replica=False, dbname='template1', username=None, max_attem
                         If False, waits for the instance to be in primary mode. Default is False.
                max_attempts:
         """
-        if not username:
-            username = default_username()
         self.start()
 
         if replica:
@@ -694,7 +693,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,
+                              username=username or self.username,
                               suppress={InternalError,
                                         QueryException,
                                         ProgrammingError,
@@ -967,15 +966,13 @@ def psql(self,
             >>> psql(query='select 3', ON_ERROR_STOP=1)
         """
 
-        # Set default arguments
         dbname = dbname or default_dbname()
-        username = username or default_username()
 
         psql_params = [
             self._get_bin_path("psql"),
             "-p", str(self.port),
             "-h", self.host,
-            "-U", username,
+            "-U", username or self.username,
             "-X",  # no .psqlrc
             "-A",  # unaligned output
             "-t",  # print rows only
@@ -1087,18 +1084,15 @@ def tmpfile():
                 fname = self.os_ops.mkstemp(prefix=TMP_DUMP)
             return fname
 
-        # Set default arguments
-        dbname = dbname or default_dbname()
-        username = username or default_username()
         filename = filename or tmpfile()
 
         _params = [
             self._get_bin_path("pg_dump"),
             "-p", str(self.port),
             "-h", self.host,
             "-f", filename,
-            "-U", username,
-            "-d", dbname,
+            "-U", username or self.username,
+            "-d", dbname or default_dbname(),
             "-F", format.value
         ]  # yapf: disable
 
@@ -1118,7 +1112,7 @@ def restore(self, filename, dbname=None, username=None):
 
         # Set default arguments
         dbname = dbname or default_dbname()
-        username = username or default_username()
+        username = username or self.username
 
         _params = [
             self._get_bin_path("pg_restore"),
@@ -1388,15 +1382,13 @@ def pgbench(self,
         if options is None:
             options = []
 
-        # Set default arguments
         dbname = dbname or default_dbname()
-        username = username or default_username()
 
         _params = [
             self._get_bin_path("pgbench"),
             "-p", str(self.port),
             "-h", self.host,
-            "-U", username,
+            "-U", username or self.username
         ] + options  # yapf: disable
 
         # should be the last one
@@ -1463,15 +1455,13 @@ def pgbench_run(self, dbname=None, username=None, options=[], **kwargs):
             >>> pgbench_run(time=10)
         """
 
-        # Set default arguments
         dbname = dbname or default_dbname()
-        username = username or default_username()
 
         _params = [
             self._get_bin_path("pgbench"),
             "-p", str(self.port),
             "-h", self.host,
-            "-U", username,
+            "-U", username or self.username
         ] + options  # yapf: disable
 
         for key, value in iteritems(kwargs):
diff --git a/testgres/operations/local_ops.py b/testgres/operations/local_ops.py
@@ -130,10 +130,6 @@ def set_env(self, var_name, var_val):
         # Check if the directory is already in PATH
         os.environ[var_name] = var_val
 
-    # Get environment variables
-    def get_user(self):
-        return self.username or getpass.getuser()
-
     def get_name(self):
         return os.name
 
diff --git a/testgres/operations/os_ops.py b/testgres/operations/os_ops.py
@@ -1,3 +1,4 @@
+import getpass
 import locale
 
 try:
@@ -44,9 +45,8 @@ def set_env(self, var_name, var_val):
         # Check if the directory is already in PATH
         raise NotImplementedError()
 
-    # Get environment variables
     def get_user(self):
-        raise NotImplementedError()
+        return getpass.getuser()
 
     def get_name(self):
         raise NotImplementedError()
