# Copyright (c) 2014 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
import logging
import time

from autotest_lib.client.common_lib import error
from autotest_lib.server import test


class platform_ServoPowerStateController(test.test):
    """Test servo can power on and off DUT in recovery and non-recovery mode."""
    version = 1


    def initialize(self, host):
        """Initialize DUT for testing."""
        pass


    def cleanup(self):
        """Clean up DUT after servo actions."""
        if not self.host.ssh_ping():
            # Power off, then power on DUT from internal storage.
            self.controller.power_off()
            self.host.servo.switch_usbkey('off')
            self.controller.power_on(self.controller.REC_OFF)


    def confirm_dut_on(self, rec_on=False):
        """Confirm DUT is powered on, claim test failure if DUT is off.

        @param rec_on: True if DUT should boot from external USB stick as in
                       recovery mode.

        @raise TestFail: If DUT is off or DUT boot from wrong source.
        """
        if not self.host.wait_up(timeout=300):
            raise error.TestFail('power_state:%s did not turn DUT on.' %
                                 ('rec' if rec_on else 'on'))

        # Check boot source. Raise TestFail if DUT boot from wrong source.
        boot_from_usb = self.host.is_boot_from_usb()
        if boot_from_usb != rec_on:
            boot_source = ('USB' if boot_from_usb else
                           'non-removable storage')
            raise error.TestFail('power_state:%s booted from %s.' %
                                 ('rec' if rec_on else 'on', boot_source))


    def confirm_dut_off(self, error_message):
        """Confirm DUT is off and does not turn back on after 30 seconds.

        @param error_message: Error message to raise if DUT stays on.
        @raise TestFail: If DUT stays on.
        """
        if not self.host.ping_wait_down(timeout=10):
            raise error.TestFail(error_message)

        if self.host.ping_wait_up(timeout=30):
            raise error.TestFail('%s. %s' % (error_message, 'DUT turns back on'
                                             'after it is turned off.'))


    def test_with_usb_plugged_in(self):
        """Run test when USB stick is plugged in servo.
        """
        logging.info('Power off DUT')
        self.controller.power_off()
        self.confirm_dut_off('power_state:off did not turn DUT off.')

        # Power DUT on in non-recovery mode with USB stick plugged in.
        # DUT shall boot from internal storage.
        logging.info('Power on DUT in non-recovery mode.')
        self.controller.power_on(self.controller.REC_OFF)
        self.confirm_dut_on()

        logging.info('Power off DUT which is up in non-recovery mode.')
        self.controller.power_off()
        self.confirm_dut_off('power_state:off failed after boot from internal '
                             'storage.')

        logging.info('Power DUT on in recovery mode, DUT shall boot from USB.')
        self.host.servo.switch_usbkey('off')
        self.controller.power_on(self.controller.REC_ON)
        self.host.servo.switch_usbkey('dut')
        time.sleep(30)
        self.confirm_dut_on(rec_on=True)

        logging.info('Power off DUT which is up in recovery mode.')
        self.controller.power_off()
        self.confirm_dut_off('power_state:off failed after boot from external '
                             'USB stick.')

        logging.info('Power on in non-recovery mode.')
        self.controller.power_on(self.controller.REC_OFF)
        self.confirm_dut_on(rec_on=False)


    def test_with_usb_unplugged(self):
        """Run test when USB stick is not plugged in servo.
        """
        # Power off DUT regardless its current status.
        logging.info('Power off DUT.')
        self.controller.power_off()
        self.confirm_dut_off('power_state:off did not turn DUT off.')

        # Try to power off the DUT again, make sure the DUT stays off.
        logging.info('Power off DUT which is already off.')
        self.controller.power_off()
        self.confirm_dut_off('power_state:off turned DUT on.')

        # USB stick should be unplugged before the test.
        self.host.servo.switch_usbkey('off')

        # DUT should be off before calling power_on(REC_ON)
        logging.info('Power on DUT in recovery mode with USB unplugged.')
        self.controller.power_on(self.controller.REC_ON)
        self.confirm_dut_off("power_state:rec didn't stay at recovery screen.")

        logging.info('Power off DUT in recovery mode with USB unplugged.')
        self.controller.power_off()
        self.confirm_dut_off('power_state:off failed at recovery screen.')

        logging.info('Power on in non-recovery mode.')
        self.controller.power_on(self.controller.REC_OFF)
        self.confirm_dut_on(rec_on=False)

        logging.info('Power DUT off and on without delay. DUT should be on '
                     'after power_on is completed.')
        self.controller.power_off()
        self.controller.power_on(self.controller.REC_OFF)
        self.confirm_dut_on(rec_on=False)

        logging.info('Power DUT off and reset. DUT should be on after reset '
                     'is completed.')
        self.controller.power_off()
        self.controller.reset()
        self.confirm_dut_on(rec_on=False)

        logging.info('Reset DUT when it\'s on. DUT should be on after reset '
                     'is completed.')
        boot_id = self.host.get_boot_id()
        self.controller.reset()
        self.confirm_dut_on(rec_on=False)
        new_boot_id = self.host.get_boot_id()
        if not new_boot_id or boot_id == new_boot_id:
            raise error.TestFail('power_state:reset failed to reboot DUT.')


    def run_once(self, host, usb_available=True):
        """Run the test.

        @param host: host object of tested DUT.
        @param usb_plugged_in: True if USB stick is plugged in servo.
        """
        self.host = host
        self.controller = host.servo.get_power_state_controller()

        self.test_with_usb_unplugged()
        if usb_available:
            self.host.servo.switch_usbkey('dut')
            self.test_with_usb_plugged_in()
            self.host.servo.switch_usbkey('off')