#!/usr/bin/python
# Copyright (c) 2010 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.

'''Unittest wrapper for chromeos_interface.py.'''

import os
import subprocess
import sys
import tempfile
import unittest

import chromeos_interface

def zero_recoverysw_count(vector_map):
    '''Calculate number of elements where recoverysw is zero in recovery mode'''
    zero_count = 0
    for key, values in vector_map.iteritems():
        if values['mainfw_type'] != 'recovery':
            continue
        if 'recoverysw_boot' in values and values['recoverysw_boot'] == '0':
            zero_count += 1
    return zero_count

# A set of crossystem output values pertinent to the tests.
CROSSYSTEM_SAMPLE = {
    'devsw_boot': '0',
    'recoverysw_boot': '0',
    'recovery_reason': '0',
    'tried_fwb': '0',
    'fwid': 'Alex.03.61.0734.0055G1.0020',
    'mainfw_act': 'A',
    'mainfw_type': 'normal',
    'ecfw_act': 'RW',
    'fwb_tries': '0',
}

class StdOut(object):
    '''
    An object to simulate the stdout file returned by subprocess.Popen().

    @text - a potentially multiline string to represent the expected output of
            the command.
    '''

    def __init__(self, text):
        (fd, self.tmp_file) = tempfile.mkstemp()
        os.write(fd, text)
        os.close(fd)
        self.stdout = open(self.tmp_file, 'r')

    def __del__(self):
        self.stdout.close()
        os.remove(self.tmp_file)


class CrosIfMock(chromeos_interface.ChromeOSInterface):
    '''A wrapper to allow testing of ChromeOSInterface on a host.'''

    def __init__(self):
        self.cs_values = CROSSYSTEM_SAMPLE
        super(CrosIfMock, self).__init__(True)
        self.init()
        self.test_write_command = ''
        self.write_command_seen = False

    def run_shell_command(self, cmd):
        if cmd.startswith('crossystem') and not '=' in cmd:
            if ' ' in cmd:
                param = cmd.split()[1]
                return StdOut(self.cs_values[param])
            else:
                return StdOut('\n'.join('%s=%s' % (x, y) for (
                            x, y) in CROSSYSTEM_SAMPLE.iteritems()))

        if not self.target_hosted():
            if cmd == 'rootdev -s':
                return StdOut('/dev/sda3')

        if cmd == self.test_write_command:
            self.test_write_command = ''
            self.write_command_seen = True
            return StdOut('')

        return super(CrosIfMock, self).run_shell_command(cmd)


CHROS_IF = CrosIfMock()

class TestShellCommands(unittest.TestCase):

    TEST_CMD = 'ls /etc'

    def setUp(self):
        '''Cache text expected to be generated by the command.'''
        self.expected_text = subprocess.Popen(
            self.TEST_CMD, shell=True,
            stdout=subprocess.PIPE).stdout.read().strip()

    def test_run_shell_command(self):
        '''Verify that the process provides correct stdout.'''
        proc = CHROS_IF.run_shell_command(self.TEST_CMD)
        self.assertEqual(proc.stdout.read().strip(), self.expected_text)

    def test_run_run_command_get_output(self):
        '''Verify that the command generates correct output.'''
        live_text = CHROS_IF.run_shell_command_get_output(self.TEST_CMD)
        self.assertEqual(live_text, self.expected_text.split('\n'))

    def test_boot_state_vector(self):
        '''Verify the format (not the contents) of the boot vector.'''
        bv = CHROS_IF.boot_state_vector()
        self.assertEqual(bv, '1:1:1:0:3')
        saved_mainfw_type = CROSSYSTEM_SAMPLE['mainfw_type']
        del(CROSSYSTEM_SAMPLE['mainfw_type'])
        self.assertRaises(KeyError, CHROS_IF.boot_state_vector)
        CROSSYSTEM_SAMPLE['mainfw_type'] = saved_mainfw_type
        CROSSYSTEM_SAMPLE['mainfw_act'] = 'X'
        try:
            CHROS_IF.boot_state_vector()
        except chromeos_interface.ChromeOSInterfaceError, e:
            dump = '\n'.join('%s=%s' % (x, y) for (
                    x, y) in CROSSYSTEM_SAMPLE.iteritems())
            self.assertEqual(e[0], dump)

    def test_crossystem_set(self):
        '''Test if appropriate crossystem command is invoked on assignments.'''
        CHROS_IF.test_write_command = 'crossystem "xxx=zzz"'
        CHROS_IF.write_command_seen = False
        CHROS_IF.cs.xxx = 'zzz'
        self.assertTrue(CHROS_IF.write_command_seen)

    def test_run_bad_shell_command(self):
        '''Test if run_shell_command() raises exception on errors.'''

        saved = sys.stdout
        sys.stdout = open('/dev/null', 'w')
        try:
            CHROS_IF.run_shell_command('ls .')
            CHROS_IF.run_shell_command('nonexisting_command')
        except chromeos_interface.ChromeOSInterfaceError, e:
            self.assertEqual(e[0], 'command nonexisting_command failed')
        finally:
            sys.stdout = saved

    def test_mario_workwaround(self):
        '''Test if the vector map is adjusted on Mario.'''

        chros_if = CrosIfMock()
        zero_values_count = zero_recoverysw_count(
            chros_if.cs.VECTOR_MAPS[0])
        self.assertNotEqual(zero_values_count, 0)

        CROSSYSTEM_SAMPLE['fwid'] = CROSSYSTEM_SAMPLE['fwid'].replace(
            'Alex', 'Mario')
        chros_if = CrosIfMock()
        zero_values_count = zero_recoverysw_count(
            chros_if.cs.VECTOR_MAPS[0])
        self.assertEqual(zero_values_count, 0)

if __name__ == '__main__':
    unittest.main()