debugging/brillo_fastboot_test.py

#!/usr/bin/env python
#
# Copyright (C) 2015 The Android Open Source Project
#
# 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 to run `fastboot` tests for Brillo.

This script can run some basic fastboot tests and/or run the `provision`
script any number of times. If the fastboot tests fail this script will
exit without attempting to provision.

The purpose of being able to `provision` multiple times is for potential
long-term stress-testing, and should generally only be necessary before
releasing changes that could affect the bootloader.

Requirements:
    1. Source envsetup.sh and lunch your target.
    2. Build using the in-tree workflow.
    3. Put the device in fastboot mode (e.g. `adb reboot bootloader`).

Example usage:
  # Fastboot tests and provision once.
  $ brillo_fastboot_test.py

  # Fastboot tests only.
  $ brillo_fastboot_test.py --skip-provision

  # Provision once only.
  $ brillo_fastboot_test.py --skip-fastboot

  # Fastboot and provision 5 times.
  $ brillo_fastboot_test.py --multi-provision 5
"""


import argparse
import os
import subprocess
import sys
import time
import unittest

from fastboot import device


class FastbootTest(unittest.TestCase):
    """Tests basic fastboot functionality."""

    # A bootloader required variable we can query.
    _REQUIRED_VAR = 'product'

    # A required partition we can flash.
    _REQUIRED_PARTITION = 'userdata'

    def setUp(self):
        self.fastboot = device.FastbootDevice()

    def test_getvar(self):
        """Tests `fastboot getvar`."""
        self.assertEqual('', self.fastboot.getvar('NotARealVariable'))
        self.assertNotEqual('', self.fastboot.getvar(self._REQUIRED_VAR))

    def test_getvar_all(self):
        """Tests `fastboot getvar all`.

        This isn't strictly necessary so the test will just be skipped
        if it doesn't pass.
        """
        all_vars = self.fastboot.getvar_all()
        if not all_vars:
            self.skipTest('Device has not implemented "getvar all"')
        self.assertEqual(self.fastboot.getvar(self._REQUIRED_VAR),
                         all_vars[self._REQUIRED_VAR])

    def test_flash(self):
        """Tests `fastboot flash` to a specific partition."""
        self.fastboot.flash(partition=self._REQUIRED_PARTITION)


def provision_test(count):
    """Tests that the provision script behaves properly.

    This is easier to implement ourselves rather than subclass
    unittest.TestCase, for a few reasons:
        * unittest has no good way to parameterize tests.
        * We want to stop testing as soon as a provisioning fails.

    Args:
        count: Number of times to provision the device.

    Raises:
        RuntimeError: environment was not set up properly.
        subprocess.CalledProcessError: provision returned non-zero.
    """
    fastboot = device.FastbootDevice()

    # provision-device lives in ANDROID_PRODUCT_OUT dir.
    if 'ANDROID_PRODUCT_OUT' not in os.environ:
        raise RuntimeError('You must source envsetup.sh and lunch a'
                           ' target first')
    script_path = os.path.join(os.environ['ANDROID_PRODUCT_OUT'],
                               'provision-device')
    if not os.path.isfile(script_path):
        raise RuntimeError('Could not find provision script at {}. Make sure to'
                           ' do a full tree build first.'.format(script_path))

    for i in range(count):
        print('[{}] Provisioning device ({}/{})'.format(time.strftime('%c'),
                                                        i + 1, count))
        subprocess.check_output(script_path, stderr=subprocess.STDOUT)
        fastboot.reboot(bootloader=True)
    print('Provisioning successful')


def run_tests(fastboot=True, provision=1):
    """Run the tests.

    Args:
        fastboot: Run fastboot tests.
        provision: How many times to run the provision tests.

    Returns:
        True if all tests passed.
    """
    # Fastboot constructor throws if there's not exactly one fastboot device.
    # Do this now to provide a single error message in this case rather than a
    # bunch of failed tests.
    try:
        device.FastbootDevice()
    except device.FastbootError as e:
        print 'Error: ' + str(e)
        print ('Make sure to source envsetup.sh, lunch, and build first, and'
               ' that exactly one fastboot device is connected.')
        return False

    if fastboot:
        test = unittest.TestLoader().loadTestsFromTestCase(FastbootTest)
        # Run with maximum print verbosity.
        result = unittest.TextTestRunner(verbosity=3).run(test)
        # Don't keep going if these tests fail or we risk trying to provision
        # with a broken fastboot.
        if not result.wasSuccessful():
            return False

    if provision > 0:
        provision_test(provision)

    return True


def main():
    """Parses command line and runs the appropriate tests."""
    parser = argparse.ArgumentParser(
            description=__doc__,
            formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('-F', '--skip-fastboot', action='store_true',
                        help='Skip low-level fastboot tests')
    parser.add_argument('-P', '--skip-provision', action='store_true',
                        help='Skip provision tests')
    parser.add_argument('-m', '--multi-provision', metavar='COUNT', type=int,
                        default=1, help='Run the provision tests COUNT times')
    options = parser.parse_args()

    provision = (0 if options.skip_provision else options.multi_provision)
    if run_tests(fastboot=not options.skip_fastboot, provision=provision):
        return 0
    return 1


if __name__ == '__main__':
    sys.exit(main())