tests/flash/FlashTest.py

#!/usr/bin/env python3
#
# Copyright 2022 The Fuchsia Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Script for to flash Fuchsia devices and reports the DUT's version of Fuchsia in
the Sponge test result properties. Uses the built in flashing tool for
fuchsia_devices.
"""

import logging

from mobly import asserts, signals, test_runner

from antlion import base_test
from antlion.controllers import fuchsia_device, pdu
from antlion.controllers.fuchsia_device import FuchsiaDevice
from antlion.controllers.pdu import PduDevice
from antlion.utils import get_device

MAX_FLASH_ATTEMPTS = 3


class FlashTest(base_test.AntlionBaseTest):
    def setup_class(self):
        self.log = logging.getLogger()
        self.fuchsia_devices: list[FuchsiaDevice] = self.register_controller(
            fuchsia_device
        )
        self.pdu_devices: list[PduDevice] = self.register_controller(pdu)
        self.failed_to_get_device_info = False

    def teardown_class(self):
        # Verify that FlashTest successfully reported the DUT version. This is
        # working around a flaw in ACTS where signals.TestAbortAll does not
        # report any errors.
        #
        # TODO(http://b/253515812): This has been fixed in Mobly already. Remove
        # teardown_class and change "TestError" to "abort_all" in
        # test_flash_devices once we move to Mobly.
        if self.failed_to_get_device_info:
            asserts.abort_all("Failed to get DUT device information")

        return super().teardown_class()

    def test_flash_devices(self) -> None:
        """Flashes a Fuchsia device for testing.

        This method calls the fuchsia_device reboot() with 'flash' argument.
        This kicks off a flash, not pave, of the fuchsia device. It also soft
        reboots the device. On error it will attempt to reflash up to
        MAX_FLASH_ATTEMPTS hard rebooting inbetween each attempt.
        """
        for device in self.fuchsia_devices:
            flash_counter = 0
            while True:
                try:
                    device.reboot(
                        reboot_type="flash", use_ssh=True, unreachable_timeout=120
                    )
                    self.log.info(f"{device.orig_ip} has been flashed.")
                    break
                except Exception as err:
                    self.log.error(
                        f"Failed to flash {device.orig_ip} with error:\n{err}"
                    )

                    if not device.device_pdu_config:
                        asserts.abort_all(
                            f"Failed to flash {device.orig_ip} and no PDU"
                            "available for hard reboot"
                        )

                    flash_counter = flash_counter + 1
                    if flash_counter == MAX_FLASH_ATTEMPTS:
                        asserts.abort_all(
                            f"Failed to flash {device.orig_ip} after"
                            f"{MAX_FLASH_ATTEMPTS} attempts"
                        )

                    self.log.info(
                        f"Hard rebooting {device.orig_ip} and retrying flash."
                    )
                    device.reboot(reboot_type="hard", testbed_pdus=self.pdu_devices)

        # Report the new Fuchsia version
        try:
            dut = get_device(self.fuchsia_devices, "DUT")
            version = dut.version()
            device_name = dut.device_name()
            product_name = dut.product_name()

            self.record_data(
                {
                    "sponge_properties": {
                        "DUT_VERSION": version,
                        "DUT_NAME": device_name,
                        "DUT_PRODUCT": product_name,
                    },
                }
            )

            self.log.info(f"DUT version: {version}")
            self.log.info(f"DUT name: {device_name}")
            self.log.info(f"DUT product: {product_name}")
        except Exception as e:
            self.failed_to_get_device_info = True
            raise signals.TestError(f"Failed to get DUT device information: {e}") from e


if __name__ == "__main__":
    test_runner.main()