run_device_tests.py: Refactor main() into smaller functions

Pylint reports several errors in the main() function, which has grown
over time to include many lines:
R0914: Too many local variables (21/15) (too-many-locals)
R0912: Too many branches (18/12) (too-many-branches)
R0915: Too many statements (84/50) (too-many-statements)

Break down the function into smaller ones that are within the guidelines
recommended by the linter.

BUG=b:239100048
TEST=cros lint test/run_device_tests.py
TEST=./test/run_device_tests.py --flasher=jtrace --remote=localhost \
       --jlink_port=19020 --board dartmonkey
     => PASS
BRANCH=none

Signed-off-by: Andrea Grandi <agrandi@google.com>
Change-Id: I77d4ddd20b102894428ef6999f4f2ead53ce6d44
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/ec/+/3766499
Reviewed-by: Tom Hughes <tomhughes@chromium.org>

1 files changed, 103 insertions, 95 deletions

diff --git a/test/run_device_tests.py b/test/run_device_tests.py
index 85f092a0c8..3d7e709c0e 100755
--- a/test/run_device_tests.py
+++ b/test/run_device_tests.py
@@ -640,6 +640,70 @@ def get_test_list(config: BoardConfig, test_args) -> List[TestConfig]:
     return test_list
 
 
+def flash_and_run_test(
+    test: TestConfig, board_config: BoardConfig, args: argparse.Namespace, executor
+) -> bool:
+    """Run a single test using the test and board configuration specified"""
+    build_board = args.board
+    # If test provides this information, build image for board specified
+    # by test.
+    if test.build_board is not None:
+        build_board = test.build_board
+
+    # build test binary
+    build(test.test_name, build_board, args.compiler)
+
+    image_path = os.path.join(
+        EC_DIR, "build", build_board, test.test_name, test.test_name + ".bin"
+    )
+
+    if test.ro_image is not None:
+        try:
+            patch_image(test, image_path)
+        except Exception as exception:  # pylint: disable=broad-except
+            logging.warning(
+                "An exception occurred while patching " "image: %s", exception
+            )
+            return False
+
+    # flash test binary
+    # TODO(b/158327221): First attempt to flash fails after
+    #  flash_write_protect test is run; works after second attempt.
+    flash_succeeded = False
+    for i in range(0, test.num_flash_attempts):
+        logging.debug("Flash attempt %d", i + 1)
+        if flash(image_path, args.board, args.flasher, args.remote, args.jlink_port):
+            flash_succeeded = True
+            break
+        time.sleep(1)
+
+    if not flash_succeeded:
+        logging.debug("Flashing failed after max attempts: %d", test.num_flash_attempts)
+        return False
+
+    if test.toggle_power:
+        power(board_config, on=False)
+        time.sleep(1)
+        power(board_config, on=True)
+
+    hw_write_protect(test.enable_hw_write_protect)
+
+    # run the test
+    logging.info('Running test: "%s"', test.config_name)
+
+    with ExitStack() as stack:
+        if args.remote and args.console_port:
+            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            s.connect((args.remote, args.console_port))
+            console = stack.enter_context(s.makefile(mode="rwb", buffering=0))
+        else:
+            console = stack.enter_context(
+                open(get_console(board_config), "wb+", buffering=0)
+            )
+
+        return run_test(test, console, executor=executor)
+
+
 def parse_remote_arg(remote: str) -> str:
     """Convert the 'remote' input argument to IP address, if available."""
     if not remote:
@@ -653,6 +717,42 @@ def parse_remote_arg(remote: str) -> str:
         sys.exit(1)
 
 
+def validate_args_combination(args: argparse.Namespace):
+    """Check that the current combination of arguments is supported.
+
+    Not all combinations of command line arguments are valid or currently
+    supported. If tests can't be executed, print and error message and exit.
+    """
+    if args.jlink_port and not args.flasher == JTRACE:
+        logging.error("jlink_port specified, but flasher is not set to J-Link.")
+        sys.exit(1)
+
+    if args.remote and not (args.jlink_port or args.console_port):
+        logging.error(
+            "jlink_port or console_port must be specified when using "
+            "the remote option."
+        )
+        sys.exit(1)
+
+    if (args.jlink_port or args.console_port) and not args.remote:
+        logging.error(
+            "The remote option must be specified when using the "
+            "jlink_port or console_port options."
+        )
+        sys.exit(1)
+
+    if args.remote and args.flasher == SERVO_MICRO:
+        logging.error(
+            "The remote option is not supported when flashing with servo "
+            "micro. Use J-Link instead or flash with a local servo micro."
+        )
+        sys.exit(1)
+
+    if args.board not in BOARD_CONFIGS:
+        logging.error('Unable to find a config for board: "%s"', args.board)
+        sys.exit(1)
+
+
 def main():
     """Run unit tests on device and displays the results."""
     parser = argparse.ArgumentParser()
@@ -689,6 +789,7 @@ def main():
         "--remote",
         "-n",
         help="The remote host connected to one or both of: J-Link and Servo.",
+        type=parse_remote_arg,
     )
 
     parser.add_argument(
@@ -703,108 +804,15 @@ def main():
 
     args = parser.parse_args()
     logging.basicConfig(level=args.log_level)
-
-    if args.jlink_port and not args.flasher == JTRACE:
-        logging.error("jlink_port specified, but flasher is not set to J-Link.")
-        sys.exit(1)
-
-    if args.remote and not (args.jlink_port or args.console_port):
-        logging.error(
-            "jlink_port or console_port must be specified when using "
-            "the remote option."
-        )
-        sys.exit(1)
-
-    if (args.jlink_port or args.console_port) and not args.remote:
-        logging.error(
-            "The remote option must be specified when using the "
-            "jlink_port or console_port options."
-        )
-        sys.exit(1)
-
-    if args.remote and args.flasher == SERVO_MICRO:
-        logging.error(
-            "The remote option is not supported when flashing with servo "
-            "micro. Use J-Link instead or flash with a local servo micro."
-        )
-        sys.exit(1)
-
-    if args.board not in BOARD_CONFIGS:
-        logging.error('Unable to find a config for board: "%s"', args.board)
-        sys.exit(1)
+    validate_args_combination(args)
 
     board_config = BOARD_CONFIGS[args.board]
-
-    remote_ip = parse_remote_arg(args.remote)
-
     e = ThreadPoolExecutor(max_workers=1)
-
     test_list = get_test_list(board_config, args.tests)
     logging.debug("Running tests: %s", [test.config_name for test in test_list])
 
     for test in test_list:
-        build_board = args.board
-        # If test provides this information, build image for board specified
-        # by test.
-        if test.build_board is not None:
-            build_board = test.build_board
-
-        # build test binary
-        build(test.test_name, build_board, args.compiler)
-
-        image_path = os.path.join(
-            EC_DIR, "build", build_board, test.test_name, test.test_name + ".bin"
-        )
-
-        if test.ro_image is not None:
-            try:
-                patch_image(test, image_path)
-            except Exception as exception:  # pylint: disable=broad-except
-                logging.warning(
-                    "An exception occurred while patching " "image: %s", exception
-                )
-                test.passed = False
-                continue
-
-        # flash test binary
-        # TODO(b/158327221): First attempt to flash fails after
-        #  flash_write_protect test is run; works after second attempt.
-        flash_succeeded = False
-        for i in range(0, test.num_flash_attempts):
-            logging.debug("Flash attempt %d", i + 1)
-            if flash(image_path, args.board, args.flasher, remote_ip, args.jlink_port):
-                flash_succeeded = True
-                break
-            time.sleep(1)
-
-        if not flash_succeeded:
-            logging.debug(
-                "Flashing failed after max attempts: %d", test.num_flash_attempts
-            )
-            test.passed = False
-            continue
-
-        if test.toggle_power:
-            power(board_config, on=False)
-            time.sleep(1)
-            power(board_config, on=True)
-
-        hw_write_protect(test.enable_hw_write_protect)
-
-        # run the test
-        logging.info('Running test: "%s"', test.config_name)
-
-        with ExitStack() as stack:
-            if remote_ip and args.console_port:
-                s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-                s.connect((remote_ip, args.console_port))
-                console = stack.enter_context(s.makefile(mode="rwb", buffering=0))
-            else:
-                console = stack.enter_context(
-                    open(get_console(board_config), "wb+", buffering=0)
-                )
-
-            test.passed = run_test(test, console, executor=e)
+        test.passed = flash_and_run_test(test, board_config, args, e)
 
     colorama.init()
     exit_code = 0