From ba6a7a18c5ab421e3d75abb739f10aae0c34886f Mon Sep 17 00:00:00 2001 From: croy Date: Tue, 15 Apr 2025 15:01:34 -0400 Subject: [PATCH] Supporting get_pty in ssh_connection.send() Change-Id: I515a8701cfc4af56a81f089caf256618a4e03612 Signed-off-by: croy --- framework/ssh/ssh_connection.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/framework/ssh/ssh_connection.py b/framework/ssh/ssh_connection.py index a7a6cb3e..90107316 100644 --- a/framework/ssh/ssh_connection.py +++ b/framework/ssh/ssh_connection.py @@ -136,7 +136,7 @@ class SSHConnection: return is_connection_success - def send(self, cmd: str, reconnect_timeout: int = 600) -> str: + def send(self, cmd: str, reconnect_timeout: int = 600, get_pty: bool = False) -> str: """ Send a command to the SSH session and return the output. @@ -146,11 +146,12 @@ class SSHConnection: Args: cmd (str): The command to execute. reconnect_timeout (int): Time in seconds to retry the connection. + get_pty (bool): Defaults to False. Whether to request a terminal when running a 'send' command. Returns: str: The output of the command. """ - return self._execute_command("SEND", cmd, reconnect_timeout=reconnect_timeout) + return self._execute_command("SEND", cmd, reconnect_timeout=reconnect_timeout, get_pty=get_pty) def send_as_sudo(self, cmd: str, reconnect_timeout: int = 600) -> str: """ @@ -187,6 +188,7 @@ class SSHConnection: cmd: str, reconnect_timeout: int = 600, prompts: List[PromptResponse] = None, + get_pty: bool = False, ) -> str: """ Executes the given action with the given command. @@ -198,6 +200,7 @@ class SSHConnection: cmd (str): The command to run. reconnect_timeout (int): The time in seconds to wait for SSH connection. prompts (List[PromptResponse], optional): Expected prompts, if any. + get_pty (bool): Defaults to False. Whether to request a terminal when running a 'send' command. Returns: str: The output of the command. @@ -223,7 +226,7 @@ class SSHConnection: thread_manager = ThreadManager(timeout=reconnect_timeout / 10) if action == "SEND": - thread_manager.start_thread("SSH_Command", self._send, cmd) + thread_manager.start_thread("SSH_Command", self._send, cmd, get_pty) elif action == "SEND_SUDO": thread_manager.start_thread("SSH_Command", self._send_as_sudo, cmd) elif action == "SEND_EXPECT_PROMPTS": @@ -244,20 +247,21 @@ class SSHConnection: time.sleep(refresh_timeout) self.is_connected = False - def _send(self, cmd: str, timeout: int = 30) -> str: + def _send(self, cmd: str, timeout: int = 30, get_pty: bool = False) -> str: """ Sends the given command with the specified timeout. Args: cmd (str): The command to send. timeout (int): The timeout in seconds for command execution. + get_pty (bool): Defaults to False. Whether to request a terminal when running a 'send' command. Returns: str: The output of the command. """ get_logger().log_ssh(cmd) - stdin, stdout, stderr = self.client.exec_command(cmd, timeout=timeout) + stdin, stdout, stderr = self.client.exec_command(cmd, timeout=timeout, get_pty=get_pty) stdout.channel.set_combine_stderr(True) self.last_return_code = stdout.channel.recv_exit_status()