enhance based on feedback

canonical · Nov 17, 2023 · 243f4c8 · 243f4c8
1 parent b99bb18
commit 243f4c8
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 21 deletions.
diff --git a/device-connectors/src/testflinger_device_connectors/fw_devices/LVFS/LVFS.py b/device-connectors/src/testflinger_device_connectors/fw_devices/LVFS/LVFS.py
@@ -8,6 +8,7 @@
 from typing import Tuple
 from testflinger_device_connectors.fw_devices.base import AbstractDevice
 from testflinger_device_connectors import logmsg
+from enum import Enum, auto
 
 
 SSH_OPTS = "-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null"
@@ -264,25 +265,20 @@ def check_results(self) -> bool:
                 fwupd_result = False
                 continue
 
-            if new_fw["Version"] == expected_ver and update_state == 2:
+            if (
+                new_fw["Version"] == expected_ver
+                and update_state
+                == FwupdUpdateState.FWUPD_UPDATE_STATE_SUCCESS.value
+            ):
                 msg = (
                     f"[{dev_name}] firmware flashed {dev['Version']}"
                     + f" → {expected_ver}"
                 )
                 log_level = logging.INFO
             else:
-                FwupdUpdateState = [
-                    "FWUPD_UPDATE_STATE_UNKNOWN",
-                    "FWUPD_UPDATE_STATE_PENDING",
-                    "FWUPD_UPDATE_STATE_SUCCESS",
-                    "FWUPD_UPDATE_STATE_FAILED",
-                    "FWUPD_UPDATE_STATE_NEEDS_REBOOT",
-                    "FWUPD_UPDATE_STATE_FAILED_TRANSIENT",
-                    "FWUPD_UPDATE_STATE_LAST",
-                ]
                 update_err = get_results.get("UpdateError", "")
                 msg = (
-                    f"[{dev_name}] {FwupdUpdateState[update_state]}:"
+                    f"[{dev_name}] {FwupdUpdateState(update_state).name}:"
                     + f" {update_err}"
                 )
                 log_level = logging.ERROR
@@ -306,13 +302,20 @@ def check_connectable(self, timeout: int):
         timeout_start = time.time()
 
         while status != "0" and time.time() < timeout_start + timeout:
-            status = subprocess.check_output(
-                f"timeout 10 ssh {SSH_OPTS} {self.user}@{self.ipaddr} "
-                + "/bin/true 2>/dev/null; echo $?",
-                shell=True,
-                universal_newlines=True,
-            ).strip()
-        if status != "0" and status != "124":
+            try:
+                status = subprocess.check_output(
+                    f"ssh {SSH_OPTS} {self.user}@{self.ipaddr} "
+                    + "/bin/true 2>/dev/null; echo $?",
+                    shell=True,
+                    universal_newlines=True,
+                    timeout=10,
+                ).strip()
+            except (
+                subprocess.TimeoutExpired,
+                subprocess.CalledProcessError,
+            ):
+                pass
+        if status != "0":
             err_msg = f"Failed to SSH to {self.ipaddr} after {timeout}s"
             logmsg(logging.ERROR, err_msg)
             raise RuntimeError(err_msg)
@@ -325,3 +328,13 @@ def reboot(self):
         self.run_cmd("sudo reboot", raise_stderr=False)
         time.sleep(10)
         self.check_connectable(self.reboot_timeout)
+
+
+class FwupdUpdateState(Enum):
+    FWUPD_UPDATE_STATE_UNKNOWN = 0
+    FWUPD_UPDATE_STATE_PENDING = auto()
+    FWUPD_UPDATE_STATE_SUCCESS = auto()
+    FWUPD_UPDATE_STATE_FAILED = auto()
+    FWUPD_UPDATE_STATE_NEEDS_REBOOT = auto()
+    FWUPD_UPDATE_STATE_FAILED_TRANSIENT = auto()
+    FWUPD_UPDATE_STATE_LAST = auto()
diff --git a/device-connectors/src/testflinger_device_connectors/fw_devices/LVFS/tests/test_LVFS.py b/device-connectors/src/testflinger_device_connectors/fw_devices/LVFS/tests/test_LVFS.py
@@ -4,7 +4,10 @@
 import unittest
 import json
 from unittest.mock import patch
-from testflinger_device_connectors.fw_devices import LVFSDevice
+from testflinger_device_connectors.fw_devices import (
+    LVFSDevice,
+    FwupdUpdateState,
+)
 from testflinger_device_connectors.fw_devices.LVFS.tests import fwupd_data
 
 device_results = json.loads(fwupd_data.GET_RESULTS_RESPONSE_DATA)
@@ -90,7 +93,9 @@ def test_check_results_good(self):
             device = LVFSDevice("", "", "")
             device._parse_fwupd_raw(fwupd_data.GET_DEVICES_RESPONSE_DATA)
             device_results = json.loads(fwupd_data.GET_RESULTS_RESPONSE_DATA)
-            device_results["UpdateState"] = 2
+            device_results[
+                "UpdateState"
+            ] = FwupdUpdateState.FWUPD_UPDATE_STATE_SUCCESS.value
             device_results["Releases"][0]["Version"] = "2.90"
             device.fw_info[2]["targetVersion"] = "2.90"
             self.assertTrue(device.check_results())

diff --git a/device-connectors/src/testflinger_device_connectors/fw_devices/__init__.py b/device-connectors/src/testflinger_device_connectors/fw_devices/__init__.py
@@ -1,5 +1,8 @@
 from testflinger_device_connectors.fw_devices.base import AbstractDevice
-from testflinger_device_connectors.fw_devices.LVFS.LVFS import LVFSDevice
+from testflinger_device_connectors.fw_devices.LVFS.LVFS import (
+    LVFSDevice,
+    FwupdUpdateState,
+)
 from testflinger_device_connectors.fw_devices.OEM.OEM import (
     OEMDevice,
     HPEDevice,
@@ -12,4 +15,5 @@
     "OEMDevice",
     "HPEDevice",
     "logmsg",
+    "FwupdUpdateState",
 ]