[PATCH v4 1/4] dts: add context manager for interactive shells

[jspewock@iol.unh.edu]

From: Jeremy Spewock <jspewock@iol.unh.edu>

Interactive shells are managed in a way currently where they are closed
and cleaned up at the time of garbage collection. Due to there being no
guarantee of when this garbage collection happens in Python, there is no
way to consistently know when an application will be closed without
manually closing the application yourself when you are done with it.
This doesn't cause a problem in cases where you can start another
instance of the same application multiple times on a server, but this
isn't the case for primary applications in DPDK. The introduction of
primary applications, such as testpmd, adds a need for knowing previous
instances of the application have been stopped and cleaned up before
starting a new one, which the garbage collector does not provide.

To solve this problem, a new class is added which acts as a base class
for interactive shells that enforces that instances of the
application be managed using a context manager. Using a context manager
guarantees that once you leave the scope of the block where the
application is being used for any reason, the application will be closed
immediately. This avoids the possibility of the shell not being closed
due to an exception being raised or user error. The interactive shell
class then becomes shells that can be started/stopped manually or at the
time of garbage collection rather than through a context manager.

depends-on: patch-139227 ("dts: skip test cases based on capabilities")

Signed-off-by: Jeremy Spewock <jspewock@iol.unh.edu>
---
 dts/framework/remote_session/__init__.py      |   1 +
 .../remote_session/interactive_shell.py       | 146 ++------------
 .../single_active_interactive_shell.py        | 188 ++++++++++++++++++
 dts/framework/remote_session/testpmd_shell.py |   9 +-
 dts/framework/testbed_model/os_session.py     |   4 +-
 dts/framework/testbed_model/sut_node.py       |   8 +-
 .../testbed_model/traffic_generator/scapy.py  |   2 +
 dts/tests/TestSuite_pmd_buffer_scatter.py     |  27 +--
 dts/tests/TestSuite_smoke_tests.py            |   3 +-
 9 files changed, 233 insertions(+), 155 deletions(-)
 create mode 100644 dts/framework/remote_session/single_active_interactive_shell.py

diff --git a/dts/framework/remote_session/__init__.py b/dts/framework/remote_session/__init__.py
index f18a9f2259..81839410b9 100644
--- a/dts/framework/remote_session/__init__.py
+++ b/dts/framework/remote_session/__init__.py
@@ -21,6 +21,7 @@
 from .interactive_shell import InteractiveShell
 from .python_shell import PythonShell
 from .remote_session import CommandResult, RemoteSession
+from .single_active_interactive_shell import SingleActiveInteractiveShell
 from .ssh_session import SSHSession
 from .testpmd_shell import NicCapability, TestPmdShell
 
diff --git a/dts/framework/remote_session/interactive_shell.py b/dts/framework/remote_session/interactive_shell.py
index 5cfe202e15..9d124b8245 100644
--- a/dts/framework/remote_session/interactive_shell.py
+++ b/dts/framework/remote_session/interactive_shell.py
@@ -1,149 +1,31 @@
 # SPDX-License-Identifier: BSD-3-Clause
 # Copyright(c) 2023 University of New Hampshire
 
-"""Common functionality for interactive shell handling.
+"""Interactive shell with manual stop/start functionality.
 
-The base class, :class:`InteractiveShell`, is meant to be extended by subclasses that contain
-functionality specific to that shell type. These subclasses will often modify things like
-the prompt to expect or the arguments to pass into the application, but still utilize
-the same method for sending a command and collecting output. How this output is handled however
-is often application specific. If an application needs elevated privileges to start it is expected
-that the method for gaining those privileges is provided when initializing the class.
-
-The :option:`--timeout` command line argument and the :envvar:`DTS_TIMEOUT`
-environment variable configure the timeout of getting the output from command execution.
+Provides a class that doesn't require being started/stopped using a context manager and can instead
+be started and stopped manually, or have the stopping process be handled at the time of garbage
+collection.
 """
 
-from abc import ABC
-from pathlib import PurePath
-from typing import Callable, ClassVar
-
-from paramiko import Channel, SSHClient, channel  # type: ignore[import]
-
-from framework.logger import DTSLogger
-from framework.settings import SETTINGS
+from .single_active_interactive_shell import SingleActiveInteractiveShell
 
 
-class InteractiveShell(ABC):
-    """The base class for managing interactive shells.
+class InteractiveShell(SingleActiveInteractiveShell):
+    """Adds manual start and stop functionality to interactive shells.
 
-    This class shouldn't be instantiated directly, but instead be extended. It contains
-    methods for starting interactive shells as well as sending commands to these shells
-    and collecting input until reaching a certain prompt. All interactive applications
-    will use the same SSH connection, but each will create their own channel on that
-    session.
+    Like its super-class, this class should not be instantiated directly and should instead be
+    extended. This class also provides an option for automated cleanup of the application through
+    the garbage collector.
     """
 
-    _interactive_session: SSHClient
-    _stdin: channel.ChannelStdinFile
-    _stdout: channel.ChannelFile
-    _ssh_channel: Channel
-    _logger: DTSLogger
-    _timeout: float
-    _app_args: str
-
-    #: Prompt to expect at the end of output when sending a command.
-    #: This is often overridden by subclasses.
-    _default_prompt: ClassVar[str] = ""
-
-    #: Extra characters to add to the end of every command
-    #: before sending them. This is often overridden by subclasses and is
-    #: most commonly an additional newline character.
-    _command_extra_chars: ClassVar[str] = ""
-
-    #: Path to the executable to start the interactive application.
-    path: ClassVar[PurePath]
-
-    #: Whether this application is a DPDK app. If it is, the build directory
-    #: for DPDK on the node will be prepended to the path to the executable.
-    dpdk_app: ClassVar[bool] = False
-
-    def __init__(
-        self,
-        interactive_session: SSHClient,
-        logger: DTSLogger,
-        get_privileged_command: Callable[[str], str] | None,
-        app_args: str = "",
-        timeout: float = SETTINGS.timeout,
-    ) -> None:
-        """Create an SSH channel during initialization.
-
-        Args:
-            interactive_session: The SSH session dedicated to interactive shells.
-            logger: The logger instance this session will use.
-            get_privileged_command: A method for modifying a command to allow it to use
-                elevated privileges. If :data:`None`, the application will not be started
-                with elevated privileges.
-            app_args: The command line arguments to be passed to the application on startup.
-            timeout: The timeout used for the SSH channel that is dedicated to this interactive
-                shell. This timeout is for collecting output, so if reading from the buffer
-                and no output is gathered within the timeout, an exception is thrown.
-        """
-        self._interactive_session = interactive_session
-        self._ssh_channel = self._interactive_session.invoke_shell()
-        self._stdin = self._ssh_channel.makefile_stdin("w")
-        self._stdout = self._ssh_channel.makefile("r")
-        self._ssh_channel.settimeout(timeout)
-        self._ssh_channel.set_combine_stderr(True)  # combines stdout and stderr streams
-        self._logger = logger
-        self._timeout = timeout
-        self._app_args = app_args
-        self._start_application(get_privileged_command)
-
-    def _start_application(self, get_privileged_command: Callable[[str], str] | None) -> None:
-        """Starts a new interactive application based on the path to the app.
-
-        This method is often overridden by subclasses as their process for
-        starting may look different.
-
-        Args:
-            get_privileged_command: A function (but could be any callable) that produces
-                the version of the command with elevated privileges.
-        """
-        start_command = f"{self.path} {self._app_args}"
-        if get_privileged_command is not None:
-            start_command = get_privileged_command(start_command)
-        self.send_command(start_command)
-
-    def send_command(self, command: str, prompt: str | None = None) -> str:
-        """Send `command` and get all output before the expected ending string.
-
-        Lines that expect input are not included in the stdout buffer, so they cannot
-        be used for expect.
-
-        Example:
-            If you were prompted to log into something with a username and password,
-            you cannot expect ``username:`` because it won't yet be in the stdout buffer.
-            A workaround for this could be consuming an extra newline character to force
-            the current `prompt` into the stdout buffer.
-
-        Args:
-            command: The command to send.
-            prompt: After sending the command, `send_command` will be expecting this string.
-                If :data:`None`, will use the class's default prompt.
-
-        Returns:
-            All output in the buffer before expected string.
-        """
-        self._logger.info(f"Sending: '{command}'")
-        if prompt is None:
-            prompt = self._default_prompt
-        self._stdin.write(f"{command}{self._command_extra_chars}\n")
-        self._stdin.flush()
-        out: str = ""
-        for line in self._stdout:
-            out += line
-            if prompt in line and not line.rstrip().endswith(
-                command.rstrip()
-            ):  # ignore line that sent command
-                break
-        self._logger.debug(f"Got output: {out}")
-        return out
+    def start_application(self) -> None:
+        """Start the application."""
+        self._start_application(self._get_privileged_command)
 
     def close(self) -> None:
         """Properly free all resources."""
-        self._stdin.close()
-        self._ssh_channel.close()
+        self._close()
 
     def __del__(self) -> None:
         """Make sure the session is properly closed before deleting the object."""
diff --git a/dts/framework/remote_session/single_active_interactive_shell.py b/dts/framework/remote_session/single_active_interactive_shell.py
new file mode 100644
index 0000000000..74060be8a7
--- /dev/null
+++ b/dts/framework/remote_session/single_active_interactive_shell.py
@@ -0,0 +1,188 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright(c) 2024 University of New Hampshire
+
+"""Common functionality for interactive shell handling.
+
+The base class, :class:`SingleActiveInteractiveShell`, is meant to be extended by subclasses that
+contain functionality specific to that shell type. These subclasses will often modify things like
+the prompt to expect or the arguments to pass into the application, but still utilize
+the same method for sending a command and collecting output. How this output is handled however
+is often application specific. If an application needs elevated privileges to start it is expected
+that the method for gaining those privileges is provided when initializing the class.
+
+This class is designed for applications like primary applications in DPDK where only one instance
+of the application can be running at a given time and, for this reason, is managed using a context
+manager. This context manager starts the application when you enter the context and cleans up the
+application when you exit. Using a context manager for this is useful since it allows us to ensure
+the application is cleaned up as soon as you leave the block regardless of the reason.
+
+The :option:`--timeout` command line argument and the :envvar:`DTS_TIMEOUT`
+environment variable configure the timeout of getting the output from command execution.
+"""
+
+from abc import ABC
+from pathlib import PurePath
+from typing import Callable, ClassVar
+
+from paramiko import Channel, SSHClient, channel  # type: ignore[import]
+from typing_extensions import Self
+
+from framework.exception import InteractiveCommandExecutionError
+from framework.logger import DTSLogger
+from framework.settings import SETTINGS
+
+
+class SingleActiveInteractiveShell(ABC):
+    """The base class for managing interactive shells.
+
+    This class shouldn't be instantiated directly, but instead be extended. It contains
+    methods for starting interactive shells as well as sending commands to these shells
+    and collecting input until reaching a certain prompt. All interactive applications
+    will use the same SSH connection, but each will create their own channel on that
+    session.
+
+    Interactive shells are started and stopped using a context manager. This allows for the start
+    and cleanup of the application to happen at predictable times regardless of exceptions or
+    interrupts.
+    """
+
+    _interactive_session: SSHClient
+    _stdin: channel.ChannelStdinFile
+    _stdout: channel.ChannelFile
+    _ssh_channel: Channel
+    _logger: DTSLogger
+    _timeout: float
+    _app_args: str
+    _get_privileged_command: Callable[[str], str] | None
+
+    #: Prompt to expect at the end of output when sending a command.
+    #: This is often overridden by subclasses.
+    _default_prompt: ClassVar[str] = ""
+
+    #: Extra characters to add to the end of every command
+    #: before sending them. This is often overridden by subclasses and is
+    #: most commonly an additional newline character.
+    _command_extra_chars: ClassVar[str] = ""
+
+    #: Path to the executable to start the interactive application.
+    path: ClassVar[PurePath]
+
+    #: Whether this application is a DPDK app. If it is, the build directory
+    #: for DPDK on the node will be prepended to the path to the executable.
+    dpdk_app: ClassVar[bool] = False
+
+    def __init__(
+        self,
+        interactive_session: SSHClient,
+        logger: DTSLogger,
+        get_privileged_command: Callable[[str], str] | None,
+        app_args: str = "",
+        timeout: float = SETTINGS.timeout,
+    ) -> None:
+        """Create an SSH channel during initialization.
+
+        Args:
+            interactive_session: The SSH session dedicated to interactive shells.
+            logger: The logger instance this session will use.
+            get_privileged_command: A method for modifying a command to allow it to use
+                elevated privileges. If :data:`None`, the application will not be started
+                with elevated privileges.
+            app_args: The command line arguments to be passed to the application on startup.
+            timeout: The timeout used for the SSH channel that is dedicated to this interactive
+                shell. This timeout is for collecting output, so if reading from the buffer
+                and no output is gathered within the timeout, an exception is thrown.
+        """
+        self._interactive_session = interactive_session
+        self._logger = logger
+        self._timeout = timeout
+        self._app_args = app_args
+        self._get_privileged_command = get_privileged_command
+
+    def _init_channel(self):
+        self._ssh_channel = self._interactive_session.invoke_shell()
+        self._stdin = self._ssh_channel.makefile_stdin("w")
+        self._stdout = self._ssh_channel.makefile("r")
+        self._ssh_channel.settimeout(self._timeout)
+        self._ssh_channel.set_combine_stderr(True)  # combines stdout and stderr streams
+
+    def _start_application(self, get_privileged_command: Callable[[str], str] | None) -> None:
+        """Starts a new interactive application based on the path to the app.
+
+        This method is often overridden by subclasses as their process for starting may look
+        different. A new SSH channel is initialized for the application to run on, then the
+        application is started.
+
+        Args:
+            get_privileged_command: A function (but could be any callable) that produces
+                the version of the command with elevated privileges.
+        """
+        self._init_channel()
+        start_command = f"{self.path} {self._app_args}"
+        if get_privileged_command is not None:
+            start_command = get_privileged_command(start_command)
+        self.send_command(start_command)
+
+    def send_command(self, command: str, prompt: str | None = None) -> str:
+        """Send `command` and get all output before the expected ending string.
+
+        Lines that expect input are not included in the stdout buffer, so they cannot
+        be used for expect.
+
+        Example:
+            If you were prompted to log into something with a username and password,
+            you cannot expect ``username:`` because it won't yet be in the stdout buffer.
+            A workaround for this could be consuming an extra newline character to force
+            the current `prompt` into the stdout buffer.
+
+        Args:
+            command: The command to send.
+            prompt: After sending the command, `send_command` will be expecting this string.
+                If :data:`None`, will use the class's default prompt.
+
+        Returns:
+            All output in the buffer before expected string.
+        """
+        self._logger.info(f"Sending: '{command}'")
+        if prompt is None:
+            prompt = self._default_prompt
+        self._stdin.write(f"{command}{self._command_extra_chars}\n")
+        self._stdin.flush()
+        out: str = ""
+        for line in self._stdout:
+            out += line
+            if prompt in line and not line.rstrip().endswith(
+                command.rstrip()
+            ):  # ignore line that sent command
+                break
+        self._logger.debug(f"Got output: {out}")
+        return out
+
+    def _close(self) -> None:
+        self._stdin.close()
+        self._ssh_channel.close()
+
+    def __enter__(self) -> Self:
+        """Enter the context block.
+
+        Upon entering a context block with this class, the desired behavior is to create the
+        channel for the application to use, and then start the application.
+
+        Returns:
+            Reference to the object for the application after it has been started.
+        """
+        self._start_application(self._get_privileged_command)
+        return self
+
+    def __exit__(self, *_) -> None:
+        """Exit the context block.
+
+        Upon exiting a context block with this class, we want to ensure that the instance of the
+        application is explicitly closed and properly cleaned up using its close method. Note that
+        because this method returns :data:`None` if an exception was raised within the block, it is
+        not handled and will be re-raised after the application is closed.
+
+        The desired behavior is to close the application regardless of the reason for exiting the
+        context and then recreate that reason afterwards. All method arguments are ignored for
+        this reason.
+        """
+        self._close()
diff --git a/dts/framework/remote_session/testpmd_shell.py b/dts/framework/remote_session/testpmd_shell.py
index f6783af621..17561d4dae 100644
--- a/dts/framework/remote_session/testpmd_shell.py
+++ b/dts/framework/remote_session/testpmd_shell.py
@@ -26,7 +26,7 @@
 from framework.settings import SETTINGS
 from framework.utils import StrEnum
 
-from .interactive_shell import InteractiveShell
+from .single_active_interactive_shell import SingleActiveInteractiveShell
 
 
 class TestPmdDevice(object):
@@ -82,7 +82,7 @@ class TestPmdForwardingModes(StrEnum):
     recycle_mbufs = auto()
 
 
-class TestPmdShell(InteractiveShell):
+class TestPmdShell(SingleActiveInteractiveShell):
     """Testpmd interactive shell.
 
     The testpmd shell users should never use
@@ -227,10 +227,11 @@ def set_forward_mode(self, mode: TestPmdForwardingModes, verify: bool = True):
                 f"Test pmd failed to set fwd mode to {mode.value}"
             )
 
-    def close(self) -> None:
+    def _close(self) -> None:
         """Overrides :meth:`~.interactive_shell.close`."""
+        self.stop()
         self.send_command("quit", "")
-        return super().close()
+        return super()._close()
 
     def get_capas_rxq(
         self, supported_capabilities: MutableSet, unsupported_capabilities: MutableSet
diff --git a/dts/framework/testbed_model/os_session.py b/dts/framework/testbed_model/os_session.py
index d5bf7e0401..745f0317f8 100644
--- a/dts/framework/testbed_model/os_session.py
+++ b/dts/framework/testbed_model/os_session.py
@@ -32,8 +32,8 @@
 from framework.remote_session import (
     CommandResult,
     InteractiveRemoteSession,
-    InteractiveShell,
     RemoteSession,
+    SingleActiveInteractiveShell,
     create_interactive_session,
     create_remote_session,
 )
@@ -43,7 +43,7 @@
 from .cpu import LogicalCore
 from .port import Port
 
-InteractiveShellType = TypeVar("InteractiveShellType", bound=InteractiveShell)
+InteractiveShellType = TypeVar("InteractiveShellType", bound=SingleActiveInteractiveShell)
 
 
 class OSSession(ABC):
diff --git a/dts/framework/testbed_model/sut_node.py b/dts/framework/testbed_model/sut_node.py
index 1fb536735d..7dd39fd735 100644
--- a/dts/framework/testbed_model/sut_node.py
+++ b/dts/framework/testbed_model/sut_node.py
@@ -243,10 +243,10 @@ def get_supported_capabilities(
         unsupported_capas: set[NicCapability] = set()
         self._logger.debug(f"Checking which capabilities from {capabilities} NIC are supported.")
         testpmd_shell = self.create_interactive_shell(TestPmdShell, privileged=True)
-        for capability in capabilities:
-            if capability not in supported_capas or capability not in unsupported_capas:
-                capability.value(testpmd_shell, supported_capas, unsupported_capas)
-        del testpmd_shell
+        with testpmd_shell as running_testpmd:
+            for capability in capabilities:
+                if capability not in supported_capas or capability not in unsupported_capas:
+                    capability.value(running_testpmd, supported_capas, unsupported_capas)
         return supported_capas
 
     def _set_up_build_target(self, build_target_config: BuildTargetConfiguration) -> None:
diff --git a/dts/framework/testbed_model/traffic_generator/scapy.py b/dts/framework/testbed_model/traffic_generator/scapy.py
index df3069d516..ba3a56df79 100644
--- a/dts/framework/testbed_model/traffic_generator/scapy.py
+++ b/dts/framework/testbed_model/traffic_generator/scapy.py
@@ -221,6 +221,8 @@ def __init__(self, tg_node: Node, config: ScapyTrafficGeneratorConfig):
             PythonShell, timeout=5, privileged=True
         )
 
+        self.session.start_application()
+
         # import libs in remote python console
         for import_statement in SCAPY_RPC_SERVER_IMPORTS:
             self.session.send_command(import_statement)
diff --git a/dts/tests/TestSuite_pmd_buffer_scatter.py b/dts/tests/TestSuite_pmd_buffer_scatter.py
index 3701c47408..645a66b607 100644
--- a/dts/tests/TestSuite_pmd_buffer_scatter.py
+++ b/dts/tests/TestSuite_pmd_buffer_scatter.py
@@ -101,7 +101,7 @@ def pmd_scatter(self, mbsize: int) -> None:
         Test:
             Start testpmd and run functional test with preset mbsize.
         """
-        testpmd = self.sut_node.create_interactive_shell(
+        testpmd_shell = self.sut_node.create_interactive_shell(
             TestPmdShell,
             app_parameters=(
                 "--mbcache=200 "
@@ -112,17 +112,20 @@ def pmd_scatter(self, mbsize: int) -> None:
             ),
             privileged=True,
         )
-        testpmd.set_forward_mode(TestPmdForwardingModes.mac)
-        testpmd.start()
-
-        for offset in [-1, 0, 1, 4, 5]:
-            recv_payload = self.scatter_pktgen_send_packet(mbsize + offset)
-            self._logger.debug(f"Payload of scattered packet after forwarding: \n{recv_payload}")
-            self.verify(
-                ("58 " * 8).strip() in recv_payload,
-                f"Payload of scattered packet did not match expected payload with offset {offset}.",
-            )
-        testpmd.stop()
+        with testpmd_shell as testpmd:
+            testpmd.set_forward_mode(TestPmdForwardingModes.mac)
+            testpmd.start()
+
+            for offset in [-1, 0, 1, 4, 5]:
+                recv_payload = self.scatter_pktgen_send_packet(mbsize + offset)
+                self._logger.debug(
+                    f"Payload of scattered packet after forwarding: \n{recv_payload}"
+                )
+                self.verify(
+                    ("58 " * 8).strip() in recv_payload,
+                    "Payload of scattered packet did not match expected payload with offset "
+                    f"{offset}.",
+                )
 
     def test_scatter_mbuf_2048(self) -> None:
         """Run the :meth:`pmd_scatter` test with `mbsize` set to 2048."""
diff --git a/dts/tests/TestSuite_smoke_tests.py b/dts/tests/TestSuite_smoke_tests.py
index a553e89662..360e64eb5a 100644
--- a/dts/tests/TestSuite_smoke_tests.py
+++ b/dts/tests/TestSuite_smoke_tests.py
@@ -100,7 +100,8 @@ def test_devices_listed_in_testpmd(self) -> None:
             List all devices found in testpmd and verify the configured devices are among them.
         """
         testpmd_driver = self.sut_node.create_interactive_shell(TestPmdShell, privileged=True)
-        dev_list = [str(x) for x in testpmd_driver.get_devices()]
+        with testpmd_driver as testpmd:
+            dev_list = [str(x) for x in testpmd.get_devices()]
         for nic in self.nics_in_node:
             self.verify(
                 nic.pci in dev_list,
-- 
2.45.1