minor updates

Author: abikouo

plugins/connection/aws_ssm.py

@@ -323,21 +323,17 @@
         cmd: '/tmp/date.sh'
 """
 import getpass
-import json
 import os
-import pty
 import random
 import re
-import select
 import string
-import subprocess
 import time
 from functools import wraps
 from typing import Any
+from typing import Callable
 from typing import Dict
 from typing import Iterator
 from typing import List
-from typing import NoReturn
 from typing import Optional
 from typing import Tuple
 from typing import TypedDict
@@ -449,6 +445,18 @@ def filter_ansi(line: str, is_windows: bool) -> str:
     return line
 
 
+def _can_timeout(func: Callable[[Any], Any]) -> Callable[[Any], Any]:
+    def wrapper(self, *args, **kwargs):
+        try:
+            return func(self, *args, **kwargs)
+        except SSMProcessManagerTimeOutFailure:
+            if self._session_manager is not None:
+                self._session_manager.has_timeout = True
+            raise
+
+    return wrapper
+
+
 class CommandResult(TypedDict):
     """
     A dictionary that contains the executed command results.
@@ -474,17 +482,6 @@ class Connection(ConnectionBase):
     _s3_client = None
     MARK_LENGTH = 26
 
-    def _can_timeout(func):
-        def run(self, *args, **kwargs):
-            try:
-                func(self, *args, **kwargs)
-            except SSMProcessManagerTimeOutFailure:
-                if self._session_manager is not None:
-                    self._session_manager.has_timeout = True
-                raise
-
-        return run
-
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
 
@@ -516,7 +513,10 @@ def _connect(self) -> Any:
         self._init_clients()
         if self._session_manager is None:
             self._session_manager = SSMSessionManager(
-                self._client, self.instance_id, verbosity_display=self.verbosity_display
+                self._client,
+                self.instance_id,
+                verbosity_display=self.verbosity_display,
+                ssm_timeout=self.get_option("ssm_timeout"),
             )
         self._session_manager.start_session(
             executable=self.get_executable(),
@@ -644,11 +644,11 @@ def exec_communicate(self, cmd: str, mark_start: str, mark_begin: str, mark_end:
         win_line = ""
         begin = False
         returncode = None
-        for poll_result in self._session_manager.process_manager.poll("EXEC", cmd):
+        for poll_result in self._session_manager.poll("EXEC", cmd):
             if not poll_result:
                 continue
 
-            line = filter_ansi(self._session_manager.process_manager.stdout_readline(), self.is_windows)
+            line = filter_ansi(self._session_manager.stdout_readline(), self.is_windows)
             self.verbosity_display(4, f"EXEC stdout line: \n{line}")
 
             if not begin and self.is_windows:
@@ -669,7 +669,7 @@ def exec_communicate(self, cmd: str, mark_start: str, mark_begin: str, mark_end:
                 stdout = stdout + line
 
         # see https://github.com/pylint-dev/pylint/issues/8909)
-        return (returncode, stdout, self._session_manager.process_manager.flush_stderr())  # pylint: disable=unreachable
+        return (returncode, stdout, self._session_manager.flush_stderr())  # pylint: disable=unreachable
 
     @staticmethod
     def generate_mark() -> str:
@@ -696,10 +696,10 @@ def exec_command(self, cmd: str, in_data: bool = None, sudoable: bool = True) ->
         # Wrap command in markers accordingly for the shell used
         cmd = self._wrap_command(cmd, mark_start, mark_end)
 
-        self._session_manager.process_manager.flush_stderr()
+        self._session_manager.flush_stderr()
 
         for chunk in chunks(cmd, 1024):
-            self._session_manager.process_manager.stdin_write(to_bytes(chunk, errors="surrogate_or_strict"))
+            self._session_manager.stdin_write(to_bytes(chunk, errors="surrogate_or_strict"))
 
         return self.exec_communicate(cmd, mark_start, mark_begin, mark_end)
 

plugins/plugin_utils/ssmsessionmanager.py

@@ -16,7 +16,6 @@
 from typing import Any
 from typing import Callable
 from typing import Dict
-from typing import List
 from typing import NoReturn
 from typing import Optional
 from typing import Union
@@ -25,33 +24,32 @@
 from ansible.module_utils._text import to_bytes
 from ansible.module_utils._text import to_text
 
-"""
-- The methods related to SSM session management (e.g., start_session, exec_command, _connect) are moved into a new class named SSMSessionManager.
-- The Connection class should now delegate SSM session tasks to the SSMSessionManager.
-- The SSMSessionManager class is moved into its own file for better modularization.
-- Python type hints should be added to method signatures for better clarity and static analysis.
-- Docstrings should describe each method's purpose, inputs, outputs, and any special handling.
-- Unit tests for SSMSessionManager should be written to ensure it functions as expected and independently.
-- The code should pass all existing and new unit tests.
-"""
-
 verbosity_display_type = Callable[[int, str], None]
 
 
 class SSMProcessManagerTimeOutFailure(AnsibleConnectionFailure):
     pass
 
 
+def _create_polling_obj(fd: Any) -> Any:
+    """create polling object using select.poll, this is helpful for unit testing"""
+    poller = select.poll()
+    poller.register(fd, select.POLLIN)
+    return poller
+
+
 class ProcessManager:
-    def __init__(self, session: Any, stdout: Any, timeout: int, verbosity_display: verbosity_display_type) -> None:
+    def __init__(
+        self, instance_id: str, session: Any, stdout: Any, timeout: int, verbosity_display: verbosity_display_type
+    ) -> None:
         self._session = session
         self._stdout = stdout
         self.verbosity_display = verbosity_display
         self._timeout = timeout
         self._poller = None
-
-        self._poller = select.poll()
-        self._poller.register(self._stdout, select.POLLIN)
+        self.instance_id = instance_id
+        self._session_id = None
+        self._poller = _create_polling_obj(self._stdout)
 
     def stdin_write(self, command: Union[bytes, str]) -> None:
         self._session.stdin.write(command)
@@ -65,10 +63,8 @@ def stdout_readline(self) -> str:
     def flush_stderr(self) -> str:
         """read and return stderr with minimal blocking"""
 
-        poller = select.poll()
-        poller.register(self._session.stderr, select.POLLIN)
+        poller = _create_polling_obj(self._session.stderr)
         stderr = ""
-
         while self._session.poll() is None:
             if not poller.poll(1):
                 break
@@ -77,23 +73,27 @@ def flush_stderr(self) -> str:
             stderr = stderr + line
         return stderr
 
+    def poll_stdout(self, length: int = 1000) -> bool:
+        return bool(self._poller.poll(length))
+
     def poll(self, label: str, cmd: str) -> NoReturn:
         """Poll session to retrieve content from stdout.
 
         :param label: A label for the display (EXEC, PRE...)
         :param cmd: The command being executed
         """
         start = round(time.time())
-        yield bool(self._poller.poll(1000))
+        yield self.poll_stdout()
         while self._session.poll() is None:
             remaining = start + self._timeout - round(time.time())
             self.verbosity_display(4, f"{label} remaining: {remaining} second(s)")
             if remaining < 0:
                 raise SSMProcessManagerTimeOutFailure(f"{label} command '{cmd}' timeout on host: {self.instance_id}")
-            yield bool(self._poller.poll(1000))
+            yield self.poll_stdout()
 
     def wait_for_match(self, label: str, cmd: str, match: Union[str, Callable[[str], bool]]) -> None:
         stdout = ""
+        self.verbosity_display(4, f"{label} WAIT FOR: {match} - Command = {cmd}")
         for result in self.poll(label=label, cmd=cmd):
             if result:
                 text = self.stdout_read_text()
@@ -109,13 +109,17 @@ def wait_for_match(self, label: str, cmd: str, match: Union[str, Callable[[str],
 class SSMSessionManager:
     MARK_LENGTH = 26
 
-    def __init__(self, ssm_client: Any, instance_id: str, verbosity_display: verbosity_display_type) -> None:
+    def __init__(
+        self, ssm_client: Any, instance_id: str, ssm_timeout: int, verbosity_display: verbosity_display_type
+    ) -> None:
         self._session_id = None
         self._instance_id = instance_id
         self.verbosity_display = verbosity_display
         self._client = ssm_client
         self._has_timeout = False
         self._process_mgr = None
+        self._timeout = ssm_timeout
+        self._session = None
 
     @property
     def instance_id(self) -> str:
@@ -137,9 +141,10 @@ def session(self) -> Any:
     def stdout(self) -> Any:
         return self._stdout
 
-    @property
-    def process_manager(self) -> ProcessManager:
-        return self._process_mgr
+    def __getattr__(self, attr: str) -> Callable:
+        if self._process_mgr and hasattr(self._process_mgr, attr):
+            return getattr(self._process_mgr, attr)
+        raise AttributeError(f"class SSMSessionManager has no attribute '{attr}'")
 
     def start_session(
         self,
@@ -148,11 +153,13 @@ def start_session(
         region_name: Optional[str],
         profile_name: str,
         prepare_terminal: bool,
-        parameters: Dict[str, Any] = {},
+        parameters: Dict[str, Any] = None,
     ) -> None:
         """Start SSM Session manager session and eventually prepare terminal"""
 
         self.verbosity_display(3, f"ESTABLISH SSM CONNECTION TO: {self.instance_id}")
+        if parameters is None:
+            parameters = {}
         start_session_args = dict(Target=self.instance_id, Parameters=parameters)
         if document_name is not None:
             start_session_args["DocumentName"] = document_name
@@ -173,12 +180,18 @@ def start_session(
         self.verbosity_display(4, f"SSM COMMAND: {to_text(cmd)}")
 
         stdout_r, stdout_w = openpty()
-        session = Popen(cmd, stdin=PIPE, stdout=stdout_w, stderr=PIPE, close_fds=True, bufsize=0)
+        self._session = Popen(cmd, stdin=PIPE, stdout=stdout_w, stderr=PIPE, close_fds=True, bufsize=0)
 
         os.close(stdout_w)
         stdout = os.fdopen(stdout_r, "rb", 0)
 
-        self._process_mgr = ProcessManager(session=session, stdout=stdout, verbosity_display=self.verbosity_display)
+        self._process_mgr = ProcessManager(
+            instance_id=self.instance_id,
+            session=self._session,
+            stdout=stdout,
+            timeout=self._timeout,
+            verbosity_display=self.verbosity_display,
+        )
 
         # For non-windows Hosts: Ensure the session has started, and disable command echo and prompt.
         if prepare_terminal:
@@ -198,7 +211,9 @@ def _disable_prompt_command(self) -> None:
         self._process_mgr.stdin_write(disable_prompt_cmd)
 
         # Read output until we got expression
-        self.wait_for_match(label="DISABLE PROMPT", cmd=disable_prompt_cmd, match=disable_prompt_reply.search)
+        self._process_mgr.wait_for_match(
+            label="DISABLE PROMPT", cmd=disable_prompt_cmd, match=disable_prompt_reply.search
+        )
 
     def _disable_echo_command(self) -> None:
         """Disable echo command from the host"""
@@ -209,12 +224,15 @@ def _disable_echo_command(self) -> None:
         self._process_mgr.stdin_write(disable_echo_cmd)
 
         # Read output until we got expression
-        self.wait_for_match(label="DISABLE ECHO", cmd=disable_echo_cmd, match="stty -echo")
+        self._process_mgr.wait_for_match(label="DISABLE ECHO", cmd=disable_echo_cmd, match="stty -echo")
 
     def _prepare_terminal(self) -> None:
         """perform any one-time terminal settings"""
+        self.verbosity_display(4, "PREPARE TERMINAL")
         # Ensure SSM Session has started
-        self.wait_for_match(label="START SSM SESSION", cmd="start_session", match="Starting session with SessionId")
+        self._process_mgr.wait_for_match(
+            label="START SSM SESSION", cmd="start_session", match="Starting session with SessionId"
+        )
 
         # Disable echo command
         self._disable_echo_command()  # pylint: disable=unreachable

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

@@ -85,33 +85,6 @@ def test_initialize_ssm_client(self, mock_boto3_client):
 
         assert conn._client is mock_boto3_client
 
-    @patch("os.path.exists")
-    @patch("subprocess.Popen")
-    @patch("select.poll")
-    @patch("boto3.client")
-    def test_plugins_connection_aws_ssm_start_session(self, boto_client, s_poll, s_popen, mock_ospe):
-        pc = PlayContext()
-        new_stdin = StringIO()
-        conn = connection_loader.get("community.aws.aws_ssm", pc, new_stdin)
-        conn.get_option = MagicMock()
-        conn.get_option.side_effect = ["i1234", "executable", "abcd", "i1234"]
-        conn.host = "abc"
-        mock_ospe.return_value = True
-        boto3 = MagicMock()
-        boto3.client("ssm").return_value = MagicMock()
-        conn.start_session = MagicMock()
-        conn._session_id = MagicMock()
-        conn._session_id.return_value = "s1"
-        s_popen.return_value.stdin.write = MagicMock()
-        s_poll.return_value = MagicMock()
-        s_poll.return_value.register = MagicMock()
-        s_popen.return_value.poll = MagicMock()
-        s_popen.return_value.poll.return_value = None
-        conn._stdin_readline = MagicMock()
-        conn._stdin_readline.return_value = "abc123"
-        conn.SESSION_START = "abc"
-        conn.start_session()
-
     @patch("random.choice")
     def test_plugins_connection_aws_ssm_exec_command(self, r_choice):
         pc = PlayContext()
@@ -339,28 +312,6 @@ def test_verbosity_diplay(self, message, level, method):
             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:
     """