modified based on feedback

Author: mandar242

plugins/connection/aws_ssm.py

@@ -484,7 +484,12 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self._instance_id = None
         self._polling_obj = None
         self._has_timeout = False
-        self.terminal_manager = TerminalManager(self)
+        self.terminal_manager = TerminalManager(
+            session=self._session,
+            stdout=self._stdout,
+            poller=self._polling_obj or select.poll(),
+            verbosity_display=lambda level, msg: self.verbosity_display(level, msg),
+        )
 
         if getattr(self._shell, "SHELL_FAMILY", "") == "powershell":
             self.delegate = None

plugins/plugin_utils/terminalmanager.py

@@ -6,20 +6,27 @@
 import random
 import re
 import string
+from typing import Any
+from typing import Callable
 
 from ansible.module_utils._text import to_bytes
 from ansible.module_utils._text import to_text
 from ansible.plugins.shell.powershell import _common_args
 
 
 class TerminalManager:
-    def __init__(self, connection):
-        self.connection = connection
+    MARK_LENGTH = 26
 
-    def prepare_terminal(self) -> None:
+    def __init__(self, session: Any, stdout: Any, poller: Callable, verbosity_display: Callable) -> None:
+        self._session = session
+        self._stdout = stdout
+        self._poller = poller
+        self.verbosity_display = verbosity_display
+
+    def prepare_terminal(self, is_windows: bool) -> None:
         """perform any one-time terminal settings"""
         # No Windows setup for now
-        if self.connection.is_windows:
+        if is_windows:
             return
 
         # Ensure SSM Session has started
@@ -31,14 +38,13 @@ def prepare_terminal(self) -> None:
         # Disable prompt command
         self.disable_prompt_command()  # pylint: disable=unreachable
 
-        self.connection.verbosity_display(4, "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:
+    def wrap_command(self, cmd: str, mark_start: str, mark_end: str, is_windows: bool, shell: Any) -> str:
         """Wrap command so stdout and status can be extracted"""
-
-        if self.connection.is_windows:
+        if is_windows:
             if not cmd.startswith(" ".join(_common_args) + " -EncodedCommand"):
-                cmd = self.connection._shell._encode_script(cmd, preserve_rc=True)
+                cmd = shell._encode_script(cmd, preserve_rc=True)
             cmd = cmd + "; echo " + mark_start + "\necho " + mark_end + "\n"
         else:
             cmd = (
@@ -47,44 +53,44 @@ 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.connection.verbosity_display(4, f"wrap_command: \n'{to_text(cmd)}'")
+        self.verbosity_display(4, f"wrap_command: \n'{to_text(cmd)}'")
         return cmd
 
     def disable_echo_command(self) -> None:
         """Disable echo command from the host"""
         disable_echo_cmd = to_bytes("stty -echo\n", errors="surrogate_or_strict")
 
         # Send command
-        self.connection.verbosity_display(4, f"DISABLE ECHO Disabling Prompt: \n{disable_echo_cmd}")
-        self.connection._session.stdin.write(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.connection.poll("DISABLE ECHO", disable_echo_cmd):
+        for poll_result in self._poller("DISABLE ECHO", disable_echo_cmd):
             if poll_result:
-                stdout += to_text(self.connection._stdout.read(1024))
-                self.connection.verbosity_display(4, f"DISABLE ECHO stdout line: \n{to_bytes(stdout)}")
+                stdout += to_text(self._stdout.read(1024))
+                self.verbosity_display(4, f"DISABLE ECHO stdout line: \n{to_bytes(stdout)}")
                 match = str(stdout).find("stty -echo")
                 if match != -1:
                     break
 
     def disable_prompt_command(self) -> None:
         """Disable prompt command from the host"""
-        end_mark = "".join([random.choice(string.ascii_letters) for i in range(self.connection.MARK_LENGTH)])
+        end_mark = "".join([random.choice(string.ascii_letters) for i in range(self.MARK_LENGTH)])
         disable_prompt_cmd = to_bytes(
             "PS1='' ; bind 'set enable-bracketed-paste off'; printf '\\n%s\\n' '" + end_mark + "'\n",
             errors="surrogate_or_strict",
         )
         disable_prompt_reply = re.compile(r"\r\r\n" + re.escape(end_mark) + r"\r\r\n", re.MULTILINE)
 
         # Send command
-        self.connection.verbosity_display(4, f"DISABLE PROMPT Disabling Prompt: \n{disable_prompt_cmd}")
-        self.connection._session.stdin.write(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.connection.poll("DISABLE PROMPT", disable_prompt_cmd):
+        for poll_result in self._poller("DISABLE PROMPT", disable_prompt_cmd):
             if poll_result:
-                stdout += to_text(self.connection._stdout.read(1024))
-                self.connection.verbosity_display(4, f"DISABLE PROMPT stdout line: \n{to_bytes(stdout)}")
+                stdout += to_text(self._stdout.read(1024))
+                self.verbosity_display(4, f"DISABLE PROMPT stdout line: \n{to_bytes(stdout)}")
                 if disable_prompt_reply.search(stdout):
                     break
 
@@ -93,11 +99,11 @@ def ensure_ssm_session_has_started(self) -> None:
         until we match the following string 'Starting session with SessionId'
         """
         stdout = ""
-        for poll_result in self.connection.poll("START SSM SESSION", "start_session"):
+        for poll_result in self._poller("START SSM SESSION", "start_session"):
             if poll_result:
-                stdout += to_text(self.connection._stdout.read(1024))
-                self.connection.verbosity_display(4, f"START SSM SESSION stdout line: \n{to_bytes(stdout)}")
+                stdout += to_text(self._stdout.read(1024))
+                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.connection.verbosity_display(4, "START SSM SESSION startup output received")
+                    self.verbosity_display(4, "START SSM SESSION startup output received")
                     break

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

@@ -45,12 +45,26 @@ def poll_mock(x, y):
 def test_ensure_ssm_session_has_started(m_to_text, m_to_bytes, connection_aws_ssm, stdout_lines, timeout_failure):
     m_to_text.side_effect = str
     m_to_bytes.side_effect = str
-    connection_aws_ssm._stdout.read = MagicMock()
+    connection_aws_ssm.get_option = MagicMock(return_value=10)
+
+    mock_session = MagicMock()
+    mock_session.poll.return_value = None
+    connection_aws_ssm._session = mock_session
 
+    connection_aws_ssm._stdout = MagicMock()
     connection_aws_ssm._stdout.read.side_effect = stdout_lines
 
+    poll_mock.results = [True for i in range(len(stdout_lines))]
+    connection_aws_ssm.poll = MagicMock()
+    connection_aws_ssm.poll.side_effect = poll_mock
+
     if not hasattr(connection_aws_ssm, "terminal_manager"):
-        connection_aws_ssm.terminal_manager = TerminalManager(connection_aws_ssm)
+        connection_aws_ssm.terminal_manager = TerminalManager(
+            session=connection_aws_ssm._session,
+            stdout=connection_aws_ssm._stdout,
+            poller=connection_aws_ssm.poll,
+            verbosity_display=lambda lvl, msg: None,
+        )
 
     poll_mock.results = [True for i in range(len(stdout_lines))]
     connection_aws_ssm.poll = MagicMock()
@@ -86,7 +100,12 @@ def test_disable_echo_command(m_to_text, m_to_bytes, connection_aws_ssm, stdout_
     connection_aws_ssm.poll.side_effect = poll_mock
 
     if not hasattr(connection_aws_ssm, "terminal_manager"):
-        connection_aws_ssm.terminal_manager = TerminalManager(connection_aws_ssm)
+        connection_aws_ssm.terminal_manager = TerminalManager(
+            session=connection_aws_ssm._session,
+            stdout=connection_aws_ssm._stdout,
+            poller=connection_aws_ssm.poll,
+            verbosity_display=lambda lvl, msg: None,
+        )
 
     if timeout_failure:
         with pytest.raises(TimeoutError):
@@ -110,7 +129,12 @@ def test_disable_prompt_command(m_to_text, m_to_bytes, m_random, connection_aws_
     connection_aws_ssm.poll.side_effect = poll_mock
 
     if not hasattr(connection_aws_ssm, "terminal_manager"):
-        connection_aws_ssm.terminal_manager = TerminalManager(connection_aws_ssm)
+        connection_aws_ssm.terminal_manager = TerminalManager(
+            session=connection_aws_ssm._session,
+            stdout=connection_aws_ssm._stdout,
+            poller=connection_aws_ssm.poll,
+            verbosity_display=lambda lvl, msg: None,
+        )
 
     m_random.choice = MagicMock()
     m_random.choice.side_effect = lambda x: "a"