Merge pull request #193 from canonical/use-importlib

Use importlib instead of imp

device-connectors/src/testflinger_device_connectors/cmd.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-# Copyright (C) 2023 Canonical
+# Copyright (C) 2023-2024 Canonical
 #
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU General Public License as published by
@@ -19,33 +19,42 @@
 
 import argparse
 import logging
+import sys
 
-from testflinger_device_connectors.devices import load_devices
+from testflinger_device_connectors.devices import (
+    DEVICE_CONNECTORS,
+    get_device_stage_func,
+)
 
 logger = logging.getLogger()
 
 
-def main():
+STAGES = (
+    "provision",
+    "firmware_update",
+    "runtest",
+    "allocate",
+    "reserve",
+    "cleanup",
+)
+
+
+def get_args(argv=None):
     """main command function for testflinger-device-connectors"""
-    devices = load_devices()
     parser = argparse.ArgumentParser()
 
     # First add a subcommand for each supported device type
-    dev_parser = parser.add_subparsers()
-    for dev_name, dev_class in devices:
+    dev_parser = parser.add_subparsers(dest="device", required=True)
+    for dev_name in DEVICE_CONNECTORS:
         dev_subparser = dev_parser.add_parser(dev_name)
-        dev_module = dev_class()
-        # Next add the subcommands that can be used and the methods they run
-        cmd_subparser = dev_subparser.add_subparsers()
-        for cmd, func in (
-            ("provision", dev_module.provision),
-            ("firmware_update", dev_module.firmware_update),
-            ("runtest", dev_module.runtest),
-            ("allocate", dev_module.allocate),
-            ("reserve", dev_module.reserve),
-            ("cleanup", dev_module.cleanup),
-        ):
-            cmd_parser = cmd_subparser.add_parser(cmd)
+
+        # Next add the subcommands that can be used
+        cmd_subparser = dev_subparser.add_subparsers(
+            dest="stage", required=True
+        )
+
+        for stage in STAGES:
+            cmd_parser = cmd_subparser.add_parser(stage)
             cmd_parser.add_argument(
                 "-c",
                 "--config",
@@ -55,6 +64,14 @@ def main():
             cmd_parser.add_argument(
                 "job_data", help="Testflinger json data file"
             )
-            cmd_parser.set_defaults(func=func)
-    args = parser.parse_args()
-    raise SystemExit(args.func(args))
+
+    return parser.parse_args(argv)
+
+
+def main():
+    """
+    Dynamically load the selected module and call the selected method
+    """
+    args = get_args()
+    func = get_device_stage_func(args.device, args.stage)
+    sys.exit(func(args))

device-connectors/src/testflinger_device_connectors/devices/__init__.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2015 Canonical
+# Copyright (C) 2015-2024 Canonical
 #
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU General Public License as published by
@@ -12,7 +12,6 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>
 
-import imp
 import json
 import logging
 import multiprocessing
@@ -22,6 +21,8 @@
 import subprocess
 import time
 from datetime import datetime, timedelta
+from importlib import import_module
+from typing import Callable
 
 import yaml
 
@@ -31,6 +32,23 @@
 )
 
 
+DEVICE_CONNECTORS = (
+    "cm3",
+    "dell_oemscript",
+    "dragonboard",
+    "hp_oemscript",
+    "iotscript",
+    "lenovo_oemscript",
+    "maas2",
+    "multi",
+    "muxpi",
+    "netboot",
+    "noprovision",
+    "oemrecovery",
+    "oemscript",
+)
+
+
 class ProvisioningError(Exception):
     pass
 
@@ -329,21 +347,12 @@ def wrapper(*args, **kwargs):
     return _wrapper
 
 
-def load_devices():
-    devices = []
-    device_path = os.path.dirname(os.path.realpath(__file__))
-    devs = [
-        os.path.join(device_path, device)
-        for device in os.listdir(device_path)
-        if os.path.isdir(os.path.join(device_path, device))
-    ]
-    for device in devs:
-        if "__pycache__" in device:
-            continue
-        module = imp.load_source("module", os.path.join(device, "__init__.py"))
-        devices.append((module.device_name, module.DeviceConnector))
-    return tuple(devices)
-
-
-if __name__ == "__main__":
-    load_devices()
+def get_device_stage_func(device: str, stage: str) -> Callable:
+    """
+    Load the selected device connector and return the selected stage method
+    """
+    module = import_module(f".{device}", package=__package__)
+    device_class = getattr(module, "DeviceConnector")
+    device_instance = device_class()
+    stage_method = getattr(device_instance, stage)
+    return stage_method

device-connectors/src/testflinger_device_connectors/devices/cm3/__init__.py

@@ -32,7 +32,6 @@
 
 
 class DeviceConnector(DefaultDevice):
-
     """Tool for provisioning baremetal with a given image."""
 
     @catch(RecoveryError, 46)

device-connectors/src/testflinger_device_connectors/devices/cm3/cm3.py

@@ -32,7 +32,6 @@
 
 
 class CM3:
-
     """Device Connector for CM3."""
 
     IMAGE_PATH_IDS = {

device-connectors/src/testflinger_device_connectors/devices/dell_oemscript/__init__.py

@@ -36,7 +36,6 @@
 
 
 class DeviceConnector(DefaultDevice):
-
     """Tool for provisioning Dell OEM devices with an oem image."""
 
     @catch(RecoveryError, 46)

device-connectors/src/testflinger_device_connectors/devices/dragonboard/__init__.py

@@ -34,7 +34,6 @@
 
 
 class DeviceConnector(DefaultDevice):
-
     """Tool for provisioning baremetal with a given image."""
 
     @catch(RecoveryError, 46)

device-connectors/src/testflinger_device_connectors/devices/dragonboard/dragonboard.py

@@ -33,7 +33,6 @@
 
 
 class Dragonboard:
-
     """Testflinger Device Connector for Dragonboard."""
 
     def __init__(self, config, job_data):

device-connectors/src/testflinger_device_connectors/devices/hp_oemscript/__init__.py

@@ -36,7 +36,6 @@
 
 
 class DeviceConnector(DefaultDevice):
-
     """Tool for provisioning HP OEM devices with an oem image."""
 
     @catch(RecoveryError, 46)

device-connectors/src/testflinger_device_connectors/devices/lenovo_oemscript/__init__.py

@@ -36,7 +36,6 @@
 
 
 class DeviceConnector(DefaultDevice):
-
     """Tool for provisioning Lenovo OEM devices with an oem image."""
 
     @catch(RecoveryError, 46)

device-connectors/src/testflinger_device_connectors/devices/maas2/__init__.py

@@ -33,7 +33,6 @@
 
 
 class DeviceConnector(DefaultDevice):
-
     """Tool for provisioning baremetal with a given image."""
 
     @catch(RecoveryError, 46)

device-connectors/src/testflinger_device_connectors/devices/maas2/maas2.py

@@ -37,7 +37,6 @@
 
 
 class Maas2:
-
     """Device Connector for Maas2."""
 
     def __init__(self, config, job_data):

device-connectors/src/testflinger_device_connectors/devices/multi/__init__.py

@@ -32,7 +32,6 @@
 
 
 class DeviceConnector(DefaultDevice):
-
     """Device Connector for provisioning multiple devices at the same time"""
 
     def init_device(self, args):

device-connectors/src/testflinger_device_connectors/devices/multi/multi.py

@@ -25,7 +25,6 @@
 
 
 class Multi:
-
     """Device Connector for multi-device"""
 
     def __init__(self, config, job_data, client):

device-connectors/src/testflinger_device_connectors/devices/muxpi/__init__.py

@@ -32,7 +32,6 @@
 
 
 class DeviceConnector(DefaultDevice):
-
     """Tool for provisioning baremetal with a given image."""
 
     @catch(RecoveryError, 46)

device-connectors/src/testflinger_device_connectors/devices/muxpi/muxpi.py

@@ -32,7 +32,6 @@
 
 
 class MuxPi:
-
     """Device Connector for MuxPi."""
 
     IMAGE_PATH_IDS = {

device-connectors/src/testflinger_device_connectors/devices/netboot/__init__.py

@@ -34,7 +34,6 @@
 
 
 class DeviceConnector(DefaultDevice):
-
     """Tool for provisioning baremetal with a given image."""
 
     @catch(RecoveryError, 46)

device-connectors/src/testflinger_device_connectors/devices/netboot/netboot.py

@@ -31,7 +31,6 @@
 
 
 class Netboot:
-
     """Testflinger Device Connector for Netboot."""
 
     def __init__(self, config):

device-connectors/src/testflinger_device_connectors/devices/noprovision/noprovision.py

@@ -29,7 +29,6 @@
 
 
 class Noprovision:
-
     """Testflinger Device Connector for Noprovision."""
 
     def __init__(self, config):

device-connectors/src/testflinger_device_connectors/devices/oemrecovery/__init__.py

@@ -33,7 +33,6 @@
 
 
 class DeviceConnector(DefaultDevice):
-
     """Tool for provisioning baremetal with a given image."""
 
     @catch(RecoveryError, 46)

device-connectors/src/testflinger_device_connectors/devices/oemrecovery/oemrecovery.py

@@ -30,7 +30,6 @@
 
 
 class OemRecovery:
-
     """Device Connector for OEM Recovery."""
 
     def __init__(self, config, job_data):

device-connectors/src/testflinger_device_connectors/devices/oemscript/__init__.py

@@ -31,7 +31,6 @@
 
 
 class DeviceConnector(DefaultDevice):
-
     """Tool for provisioning baremetal with a given image."""
 
     @catch(RecoveryError, 46)

device-connectors/src/testflinger_device_connectors/fw_devices/LVFS/LVFS.py

@@ -1,6 +1,5 @@
 """Device class for flashing firmware on device supported by LVFS-fwupd"""
 
-
 import subprocess
 import json
 import time

device-connectors/src/testflinger_device_connectors/fw_devices/LVFS/tests/test_LVFS.py

@@ -1,6 +1,5 @@
 """Test LVFSDevice"""
 
-
 import unittest
 import json
 from unittest.mock import patch
@@ -93,9 +92,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"
-            ] = FwupdUpdateState.FWUPD_UPDATE_STATE_SUCCESS.value
+            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())

device-connectors/src/testflinger_device_connectors/fw_devices/tests/test_firmware_update.py

@@ -1,6 +1,5 @@
 """Test detect_device in firmware_update.py"""
 
-
 import unittest
 import pytest
 from unittest.mock import patch

device-connectors/src/tests/test_cmd.py

@@ -0,0 +1,47 @@
+# Copyright (C) 2024 Canonical
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>
+"""Tests for the cmd argument parser"""
+
+import pytest
+from testflinger_device_connectors.cmd import get_args
+
+
+def test_good_args():
+    """Test that we can parse good arguments"""
+    argv = ["noprovision", "reserve", "-c", "config.cfg", "job_data.json"]
+
+    args = get_args(argv)
+
+    assert args.device == "noprovision"
+    assert args.stage == "reserve"
+    assert args.config == "config.cfg"
+    assert args.job_data == "job_data.json"
+
+
+def test_invalid_device():
+    """Test that an invalid device raises an exception"""
+    argv = ["INVALID", "provision", "-c", "config.cfg", "job_data.json"]
+
+    with pytest.raises(SystemExit) as err:
+        get_args(argv)
+    assert err.value.code == 2
+
+
+def test_invalid_stage():
+    """Test that an invalid stage raises an exception"""
+    argv = ["INVALID", "provision", "-c", "config.cfg", "job_data.json"]
+
+    with pytest.raises(SystemExit) as err:
+        get_args(argv)
+    assert err.value.code == 2