tests: add Serial class

Add the Serial class for serial console communication.

Signed-off-by: Marco Vedovati <[email protected]>

Author: marcov

tests/framework/microvm.py

@@ -13,6 +13,8 @@
 import os
 from queue import Queue
 import re
+import select
+import time
 from subprocess import run, PIPE
 
 from retry import retry
@@ -37,6 +39,8 @@ class Microvm:
     process.
     """
 
+    SCREEN_LOGFILE = "/tmp/screen.log"
+
     def __init__(
         self,
         resource_path,
@@ -314,14 +318,15 @@ def spawn(self):
         else:
             # Delete old screen log if any.
             try:
-                os.unlink('/tmp/screen.log')
+                os.unlink(self.SCREEN_LOGFILE)
             except OSError:
                 pass
-            # Log screen output to /tmp/screen.log.
+            # Log screen output to SCREEN_LOGFILE
             # This file will collect any output from 'screen'ed Firecracker.
-            start_cmd = 'screen -L -Logfile /tmp/screen.log '\
+            start_cmd = 'screen -L -Logfile {logfile} '\
                         '-dmS {session} {binary} {params}'
             start_cmd = start_cmd.format(
+                logfile=self.SCREEN_LOGFILE,
                 session=self._session_name,
                 binary=self._jailer_binary_path,
                 params=' '.join(jailer_param_list)
@@ -468,3 +473,63 @@ def start(self):
         """
         response = self.actions.put(action_type='InstanceStart')
         assert self._api_session.is_status_no_content(response.status_code)
+
+
+class Serial:
+    """Class for serial console communication with a Microvm."""
+
+    RX_TIMEOUT_S = 5
+
+    def __init__(self, vm):
+        """Initialize a new Serial object."""
+        self._poller = None
+        self._vm = vm
+
+    def open(self):
+        """Open a serial connection."""
+        # Open the screen log file.
+        if self._poller is not None:
+            # serial already opened
+            return
+
+        screen_log_fd = os.open(Microvm.SCREEN_LOGFILE, os.O_RDONLY)
+        self._poller = select.poll()
+        self._poller.register(screen_log_fd,
+                              select.POLLIN | select.POLLHUP)
+
+    def tx(self, input_string, end='\n'):
+        # pylint: disable=invalid-name
+        # No need to have a snake_case naming style for a single word.
+        r"""Send a string terminated by an end token (defaulting to "\n")."""
+        self._vm.serial_input(input_string + end)
+
+    def rx_char(self):
+        """Read a single character."""
+        result = self._poller.poll(0.1)
+
+        for fd, flag in result:
+            if flag & select.POLLHUP:
+                assert False, "Oh! The console vanished before test completed."
+
+            if flag & select.POLLIN:
+                output_char = str(os.read(fd, 1),
+                                  encoding='utf-8',
+                                  errors='ignore')
+                return output_char
+
+        return ''
+
+    def rx(self, token="\n"):
+        # pylint: disable=invalid-name
+        # No need to have a snake_case naming style for a single word.
+        r"""Read a string delimited by an end token (defaults to "\n")."""
+        rx_str = ''
+        start = time.time()
+        while True:
+            rx_str += self.rx_char()
+            if rx_str.endswith(token):
+                break
+            if (time.time() - start) >= self.RX_TIMEOUT_S:
+                assert False
+
+        return rx_str

tests/framework/state_machine.py

@@ -46,7 +46,7 @@ def __init__(self, match_string=''):
         MatchStaticString.__init__(self, match_string)
         print('\n*** Current test state: ', str(self), end='')
 
-    def handle_input(self, microvm, input_char):
+    def handle_input(self, serial, input_char):
         """Handle input event and return next state."""
 
     def __repr__(self):

tests/integration_tests/functional/test_serial_io.py

@@ -1,42 +1,41 @@
 # Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 """Tests scenario for the Firecracker serial console."""
-import os
 import time
-import select
+from framework.microvm import Serial
 from framework.state_machine import TestState
 
 
 class WaitLogin(TestState):
     """Initial state when we wait for the login prompt."""
 
-    def handle_input(self, microvm, input_char) -> TestState:
+    def handle_input(self, serial, input_char) -> TestState:
         """Handle input and return next state."""
         if self.match(input_char):
             # Send login username.
-            microvm.serial_input("root")
+            serial.tx("root")
             return WaitPasswordPrompt("Password:")
         return self
 
 
 class WaitPasswordPrompt(TestState):
     """Wait for the password prompt to be shown."""
 
-    def handle_input(self, microvm, input_char) -> TestState:
+    def handle_input(self, serial, input_char) -> TestState:
         """Handle input and return next state."""
         if self.match(input_char):
-            microvm.serial_input("root")
+            serial.tx("root")
             # Wait 1 second for shell
             time.sleep(1)
-            microvm.serial_input("id")
+            serial.tx("id")
             return WaitIDResult("uid=0(root) gid=0(root) groups=0(root)")
         return self
 
 
 class WaitIDResult(TestState):
     """Wait for the console to show the result of the 'id' shell command."""
 
-    def handle_input(self, microvm, input_char) -> TestState:
+    def handle_input(self, unused_serial, input_char) -> TestState:
         """Handle input and return next state."""
         if self.match(input_char):
             return TestFinished()
@@ -46,7 +45,7 @@ def handle_input(self, microvm, input_char) -> TestState:
 class TestFinished(TestState):
     """Test complete and successful."""
 
-    def handle_input(self, microvm, input_char) -> TestState:
+    def handle_input(self, unused_serial, input_char) -> TestState:
         """Return self since the test is about to end."""
         return self
 
@@ -67,29 +66,13 @@ def test_serial_console_login(test_microvm_with_ssh):
 
     microvm.start()
 
-    # Screen stdout log
-    screen_log = "/tmp/screen.log"
-
-    # Open the screen log file.
-    screen_log_fd = os.open(screen_log, os.O_RDONLY)
-    poller = select.poll()
+    serial = Serial(microvm)
+    serial.open()
 
     # Set initial state - wait for 'login:' prompt
     current_state = WaitLogin("login:")
 
-    poller.register(screen_log_fd, select.POLLIN | select.POLLHUP)
-
     while not isinstance(current_state, TestFinished):
-        result = poller.poll(0.1)
-        for fd, flag in result:
-            if flag & select.POLLIN:
-                output_char = str(os.read(fd, 1),
-                                  encoding='utf-8',
-                                  errors='ignore')
-                # [DEBUG] Uncomment to see the serial console output.
-                # print(output_char, end='')
-                current_state = current_state.handle_input(
-                                microvm, output_char)
-            elif flag & select.POLLHUP:
-                assert False, "Oh! The console vanished before test completed."
-    os.close(screen_log_fd)
+        output_char = serial.rx_char()
+        current_state = current_state.handle_input(
+            serial, output_char)