Refactor aws_ssm display and verbosity-related methods (#2264)

SUMMARY
Resolves ACA-2239
Methods like self._vvvv, self._vvv, etc. have been refactored into a single method, including code from the _display() method.
ISSUE TYPE


Feature Pull Request

COMPONENT NAME

connection/aws_ssm

Reviewed-by: Bianca Henderson <[email protected]>
Reviewed-by: Bikouo Aubin
Reviewed-by: Mike Graves <[email protected]>
Reviewed-by: Rahmanim Benny <[email protected]>

Author: beeankha

changelogs/fragments/2264-aws_ssm_refactor_verbosity_display.yaml

@@ -0,0 +1,2 @@
+minor_changes:
+  - aws_ssm - Refactor display/verbosity-related methods in aws_ssm to simplify the code and avoid repetition (https://github.com/ansible-collections/community.aws/pull/2264).

plugins/connection/aws_ssm.py

@@ -383,7 +383,7 @@ def wrapped(self, *args: Any, **kwargs: Any) -> Any:
         for attempt in range(remaining_tries):
             try:
                 return_tuple = func(self, *args, **kwargs)
-                self._vvvv(f"ssm_retry: (success) {to_text(return_tuple)}")
+                self.verbosity_display(4, f"ssm_retry: (success) {to_text(return_tuple)}")
                 break
 
             except (AnsibleConnectionFailure, Exception) as e:
@@ -400,7 +400,7 @@ def wrapped(self, *args: Any, **kwargs: Any) -> Any:
                         f"from cmd ({cmd_summary}),pausing for {pause} seconds"
                     )
 
-                self._vv(msg)
+                self.verbosity_display(2, msg)
 
                 time.sleep(pause)
 
@@ -511,7 +511,8 @@ def _init_clients(self) -> None:
         Initializes required AWS clients (SSM and S3).
         Delegates client initialization to specialized methods.
         """
-        self._vvvv("INITIALIZE BOTO3 CLIENTS")
+
+        self.verbosity_display(4, "INITIALIZE BOTO3 CLIENTS")
         profile_name = self.get_option("profile") or ""
         region_name = self.get_option("region")
 
@@ -520,7 +521,7 @@ def _init_clients(self) -> None:
 
         # Initialize S3 client
         s3_endpoint_url, s3_region_name = self.s3_manager.get_bucket_endpoint()
-        self._vvvv(f"SETUP BOTO3 CLIENTS: S3 {s3_endpoint_url}")
+        self.verbosity_display(4, f"SETUP BOTO3 CLIENTS: S3 {s3_endpoint_url}")
         self.s3_manager.initialize_client(
             region_name=s3_region_name, endpoint_url=s3_endpoint_url, profile_name=profile_name
         )
@@ -540,35 +541,36 @@ def _initialize_ssm_client(self, region_name: Optional[str], profile_name: str)
             None
         """
 
-        self._vvvv("SETUP BOTO3 CLIENTS: SSM")
+        self.verbosity_display(4, "SETUP BOTO3 CLIENTS: SSM")
         self._client = self._get_boto_client(
             "ssm",
             region_name=region_name,
             profile_name=profile_name,
         )
 
-    def _display(self, f: Any, message: str) -> None:
+    def verbosity_display(self, level: int, message: str) -> None:
+        """
+        Displays the given message depending on the verbosity level.
+
+        :param message: The message to display.
+        :param display_level: The verbosity level (1-4).
+
+        :return: None
+        """
         if self.host:
             host_args = {"host": self.host}
         else:
             host_args = {}
-        f(to_text(message), **host_args)
-
-    def _v(self, message: str) -> None:
-        self._display(display.v, message)
-
-    def _vv(self, message: str) -> None:
-        self._display(display.vv, message)
 
-    def _vvv(self, message: str) -> None:
-        self._display(display.vvv, message)
+        verbosity_level = {1: display.v, 2: display.vv, 3: display.vvv, 4: display.vvvv}
 
-    def _vvvv(self, message: str) -> None:
-        self._display(display.vvvv, message)
+        if level not in verbosity_level.keys():
+            raise AnsibleError(f"Invalid verbosity level: {level}")
+        verbosity_level[level](to_text(message), **host_args)
 
     def reset(self) -> Any:
         """start a fresh ssm session"""
-        self._vvvv("reset called on ssm connection")
+        self.verbosity_display(4, "reset called on ssm connection")
         self.close()
         return self.start_session()
 
@@ -601,13 +603,13 @@ def get_executable(self) -> str:
     def start_session(self):
         """start ssm session"""
 
-        self._vvv(f"ESTABLISH SSM CONNECTION TO: {self.instance_id}")
+        self.verbosity_display(3, f"ESTABLISH SSM CONNECTION TO: {self.instance_id}")
 
         executable = self.get_executable()
 
         self._init_clients()
 
-        self._vvvv(f"START SSM SESSION: {self.instance_id}")
+        self.verbosity_display(4, f"START SSM SESSION: {self.instance_id}")
         start_session_args = dict(Target=self.instance_id, Parameters={})
         document_name = self.get_option("ssm_document")
         if document_name is not None:
@@ -627,7 +629,7 @@ def start_session(self):
             self._client.meta.endpoint_url,
         ]
 
-        self._vvvv(f"SSM COMMAND: {to_text(cmd)}")
+        self.verbosity_display(4, f"SSM COMMAND: {to_text(cmd)}")
 
         stdout_r, stdout_w = pty.openpty()
         self._session = subprocess.Popen(
@@ -645,7 +647,7 @@ def start_session(self):
         # For non-windows Hosts: Ensure the session has started, and disable command echo and prompt.
         self._prepare_terminal()
 
-        self._vvvv(f"SSM CONNECTION ID: {self._session_id}")  # pylint: disable=unreachable
+        self.verbosity_display(4, f"SSM CONNECTION ID: {self._session_id}")  # pylint: disable=unreachable
 
         return self._session
 
@@ -671,7 +673,7 @@ def poll(self, label: str, cmd: str) -> NoReturn:
         timeout = self.get_option("ssm_timeout")
         while self._session.poll() is None:
             remaining = start + timeout - round(time.time())
-            self._vvvv(f"{label} remaining: {remaining} second(s)")
+            self.verbosity_display(4, f"{label} remaining: {remaining} second(s)")
             if remaining < 0:
                 self._has_timeout = True
                 raise AnsibleConnectionFailure(f"{label} command '{cmd}' timeout on host: {self.instance_id}")
@@ -697,7 +699,7 @@ def exec_communicate(self, cmd: str, mark_start: str, mark_begin: str, mark_end:
                 continue
 
             line = filter_ansi(self._stdout.readline(), self.is_windows)
-            self._vvvv(f"EXEC stdout line: \n{line}")
+            self.verbosity_display(4, f"EXEC stdout line: \n{line}")
 
             if not begin and self.is_windows:
                 win_line = win_line + line
@@ -710,9 +712,9 @@ def exec_communicate(self, cmd: str, mark_start: str, mark_begin: str, mark_end:
                 continue
             if begin:
                 if mark_end in line:
-                    self._vvvv(f"POST_PROCESS: \n{to_text(stdout)}")
+                    self.verbosity_display(4, f"POST_PROCESS: \n{to_text(stdout)}")
                     returncode, stdout = self._post_process(stdout, mark_begin)
-                    self._vvvv(f"POST_PROCESSED: \n{to_text(stdout)}")
+                    self.verbosity_display(4, f"POST_PROCESSED: \n{to_text(stdout)}")
                     break
                 stdout = stdout + line
 
@@ -731,7 +733,7 @@ def exec_command(self, cmd: str, in_data: bool = None, sudoable: bool = True) ->
 
         super().exec_command(cmd, in_data=in_data, sudoable=sudoable)
 
-        self._vvv(f"EXEC: {to_text(cmd)}")
+        self.verbosity_display(3, f"EXEC: {to_text(cmd)}")
 
         mark_begin = self.generate_mark()
         if self.is_windows:
@@ -758,10 +760,10 @@ def _ensure_ssm_session_has_started(self) -> None:
         for poll_result in self.poll("START SSM SESSION", "start_session"):
             if poll_result:
                 stdout += to_text(self._stdout.read(1024))
-                self._vvvv(f"START SSM SESSION stdout line: \n{to_bytes(stdout)}")
+                self.verbosity_display(4, f"START SSM SESSION stdout line: \n{to_bytes(stdout)}")
                 match = str(stdout).find("Starting session with SessionId")
                 if match != -1:
-                    self._vvvv("START SSM SESSION startup output received")
+                    self.verbosity_display(4, "START SSM SESSION startup output received")
                     break
 
     def _disable_prompt_command(self) -> None:
@@ -774,14 +776,14 @@ def _disable_prompt_command(self) -> None:
         disable_prompt_reply = re.compile(r"\r\r\n" + re.escape(end_mark) + r"\r\r\n", re.MULTILINE)
 
         # Send command
-        self._vvvv(f"DISABLE PROMPT Disabling Prompt: \n{disable_prompt_cmd}")
+        self.verbosity_display(4, f"DISABLE PROMPT Disabling Prompt: \n{disable_prompt_cmd}")
         self._session.stdin.write(disable_prompt_cmd)
 
         stdout = ""
         for poll_result in self.poll("DISABLE PROMPT", disable_prompt_cmd):
             if poll_result:
                 stdout += to_text(self._stdout.read(1024))
-                self._vvvv(f"DISABLE PROMPT stdout line: \n{to_bytes(stdout)}")
+                self.verbosity_display(4, f"DISABLE PROMPT stdout line: \n{to_bytes(stdout)}")
                 if disable_prompt_reply.search(stdout):
                     break
 
@@ -790,14 +792,14 @@ def _disable_echo_command(self) -> None:
         disable_echo_cmd = to_bytes("stty -echo\n", errors="surrogate_or_strict")
 
         # Send command
-        self._vvvv(f"DISABLE ECHO Disabling Prompt: \n{disable_echo_cmd}")
+        self.verbosity_display(4, f"DISABLE ECHO Disabling Prompt: \n{disable_echo_cmd}")
         self._session.stdin.write(disable_echo_cmd)
 
         stdout = ""
         for poll_result in self.poll("DISABLE ECHO", disable_echo_cmd):
             if poll_result:
                 stdout += to_text(self._stdout.read(1024))
-                self._vvvv(f"DISABLE ECHO stdout line: \n{to_bytes(stdout)}")
+                self.verbosity_display(4, f"DISABLE ECHO stdout line: \n{to_bytes(stdout)}")
                 match = str(stdout).find("stty -echo")
                 if match != -1:
                     break
@@ -817,7 +819,7 @@ def _prepare_terminal(self) -> None:
         # Disable prompt command
         self._disable_prompt_command()  # pylint: disable=unreachable
 
-        self._vvvv("PRE Terminal configured")  # pylint: disable=unreachable
+        self.verbosity_display(4, "PRE Terminal configured")  # pylint: disable=unreachable
 
     def _wrap_command(self, cmd: str, mark_start: str, mark_end: str) -> str:
         """Wrap command so stdout and status can be extracted"""
@@ -833,7 +835,7 @@ def _wrap_command(self, cmd: str, mark_start: str, mark_end: str) -> str:
                 f"printf '\\n%s\\n%s\\n' \"$?\" '{mark_end}';\n"
             )  # fmt: skip
 
-        self._vvvv(f"_wrap_command: \n'{to_text(cmd)}'")
+        self.verbosity_display(4, f"_wrap_command: \n'{to_text(cmd)}'")
         return cmd
 
     def _post_process(self, stdout: str, mark_begin: str) -> Tuple[str, str]:
@@ -881,7 +883,7 @@ def _flush_stderr(self, session_process) -> str:
             if not poll_stderr.poll(1):
                 break
             line = session_process.stderr.readline()
-            self._vvvv(f"stderr line: {to_text(line)}")
+            self.verbosity_display(4, f"stderr line: {to_text(line)}")
             stderr = stderr + line
 
         return stderr
@@ -1084,7 +1086,7 @@ def put_file(self, in_path: str, out_path: str) -> Tuple[int, str, str]:
 
         super().put_file(in_path, out_path)
 
-        self._vvv(f"PUT {in_path} TO {out_path}")
+        self.verbosity_display(3, f"PUT {in_path} TO {out_path}")
         if not os.path.exists(to_bytes(in_path, errors="surrogate_or_strict")):
             raise AnsibleFileNotFound(f"file or module does not exist: {in_path}")
 
@@ -1095,19 +1097,19 @@ def fetch_file(self, in_path: str, out_path: str) -> Tuple[int, str, str]:
 
         super().fetch_file(in_path, out_path)
 
-        self._vvv(f"FETCH {in_path} TO {out_path}")
+        self.verbosity_display(3, f"FETCH {in_path} TO {out_path}")
         return self._file_transport_command(in_path, out_path, "get")
 
     def close(self) -> None:
         """terminate the connection"""
         if self._session_id:
-            self._vvv(f"CLOSING SSM CONNECTION TO: {self.instance_id}")
+            self.verbosity_display(3, f"CLOSING SSM CONNECTION TO: {self.instance_id}")
             if self._has_timeout:
                 self._session.terminate()
             else:
                 cmd = b"\nexit\n"
                 self._session.communicate(cmd)
 
-            self._vvvv(f"TERMINATE SSM SESSION: {self._session_id}")
+            self.verbosity_display(4, f"TERMINATE SSM SESSION: {self._session_id}")
             self._client.terminate_session(SessionId=self._session_id)
             self._session_id = ""

plugins/plugin_utils/s3clientmanager.py

@@ -38,7 +38,7 @@ def get_bucket_endpoint(self) -> Tuple[str, str]:
         """
         region_name = self.connection.get_option("region") or "us-east-1"
         profile_name = self.connection.get_option("profile") or ""
-        self.connection._vvvv("_get_bucket_endpoint: S3 (global)")
+        self.connection.verbosity_display(4, "_get_bucket_endpoint: S3 (global)")
         tmp_s3_client = self.get_s3_client(
             region_name=region_name,
             profile_name=profile_name,
@@ -56,7 +56,7 @@ def get_bucket_endpoint(self) -> Tuple[str, str]:
             return self.connection.get_option("bucket_endpoint_url"), bucket_region
 
         # Create another client for the region the bucket lives in, so we can nab the endpoint URL
-        self.connection._vvvv(f"_get_bucket_endpoint: S3 (bucket region) - {bucket_region}")
+        self.connection.verbosity_display(4, f"_get_bucket_endpoint: S3 (bucket region) - {bucket_region}")
         s3_bucket_client = self.get_s3_client(
             region_name=bucket_region,
             profile_name=profile_name,

tests/unit/plugins/connection/aws_ssm/conftest.py

@@ -40,17 +40,8 @@ def connection_init(*args, **kwargs):
     def display_msg(msg):
         print("--- AWS SSM CONNECTION --- ", msg)
 
-    connection._v = MagicMock()
-    connection._v.side_effect = display_msg
-
-    connection._vv = MagicMock()
-    connection._vv.side_effect = display_msg
-
-    connection._vvv = MagicMock()
-    connection._vvv.side_effect = display_msg
-
-    connection._vvvv = MagicMock()
-    connection._vvvv.side_effect = display_msg
+    connection.verbosity_display = MagicMock()
+    connection.verbosity_display.side_effect = lambda level, msg: display_msg(msg)
 
     connection.get_option = MagicMock()
     return connection

tests/unit/plugins/connection/aws_ssm/test_aws_ssm.py

@@ -6,6 +6,7 @@
 
 import pytest
 
+from ansible.errors import AnsibleError
 from ansible.playbook.play_context import PlayContext
 from ansible.plugins.loader import connection_loader
 
@@ -258,7 +259,7 @@ def test_generate_mark(self):
 
     @pytest.mark.parametrize("is_windows", [False, True])
     def test_generate_commands(self, is_windows):
-        """Testing command generation on Windows systems"""
+        """Testing command generation on both Windows and non-Windows systems"""
         pc = PlayContext()
         new_stdin = StringIO()
         conn = connection_loader.get("community.aws.aws_ssm", pc, new_stdin)
@@ -307,6 +308,59 @@ def test_generate_commands(self, is_windows):
         assert isinstance(test_command_generation[0], list)
         assert isinstance(test_command_generation[0][0], dict)
 
+    @pytest.mark.parametrize(
+        "message,level,method",
+        [
+            ("test message 1", 1, "v"),
+            ("test message 2", 2, "vv"),
+            ("test message 3", 3, "vvv"),
+            ("test message 4", 4, "vvvv"),
+        ],
+    )
+    def test_verbosity_diplay(self, message, level, method):
+        """Testing verbosity levels"""
+        play_context = MagicMock()
+        play_context.shell = "sh"
+        conn = Connection(play_context)
+        conn.host = "test-host"  # Test with host set
+
+        with patch("ansible_collections.community.aws.plugins.connection.aws_ssm.display") as mock_display:
+            conn.verbosity_display(level, message)
+            # Verify the correct display method was called with expected args
+            mock_method = getattr(mock_display, method)
+            mock_method.assert_called_once_with(message, host="test-host")
+
+            # Test without host set
+            conn.host = None
+            conn.verbosity_display(1, "no host message")
+            mock_display.v.assert_called_with("no host message")
+
+            # Test exception is raised when verbosity level is not an accepted value
+            with pytest.raises(AnsibleError):
+                conn.verbosity_display("invalid value", "test message")
+
+    def test_poll_verbosity(self):
+        """Test poll method verbosity display"""
+        pc = PlayContext()
+        new_stdin = StringIO()
+        conn = connection_loader.get("community.aws.aws_ssm", pc, new_stdin)
+
+        conn._session = MagicMock()
+        conn._session.poll.return_value = None
+        conn.get_option = MagicMock(return_value=10)  # ssm_timeout
+        conn.poll_stdout = MagicMock()
+        conn.instance_id = "i-1234567890"
+        conn.host = conn.instance_id
+
+        with patch("time.time", return_value=100), patch.object(conn, "verbosity_display") as mock_display:
+            poll_gen = conn.poll("TEST", "test command")
+            # Advance generator twice to trigger the verbosity message
+            next(poll_gen)
+            next(poll_gen)
+
+            # Verify verbosity message contains remaining time
+            mock_display.assert_called_with(4, "TEST remaining: 10 second(s)")
+
 
 class TestS3ClientManager:
     """