server/hosts/servo_host_unittest.py

import mock
import unittest

import common

from autotest_lib.server.hosts import servo_host


class MockCmd(object):
    """Simple mock command with base command and results"""

    def __init__(self, cmd, exit_status, stdout):
        self.cmd = cmd
        self.stdout = stdout
        self.exit_status = exit_status


class MockHost(servo_host.ServoHost):
    """Simple host for running mock'd host commands"""

    def __init__(self, *args):
        self._mock_cmds = {c.cmd: c for c in args}
        self._init_attributes()
        self.hostname = "some_hostname"

    def run(self, command, **kwargs):
        """Finds the matching result by command value"""
        mock_cmd = self._mock_cmds[command]
        file_out = kwargs.get('stdout_tee', None)
        if file_out:
            file_out.write(mock_cmd.stdout)
        return mock_cmd


class ServoHostServoStateTestCase(unittest.TestCase):
    """Tests to verify changing the servo_state"""
    def test_return_broken_if_state_not_defined(self):
        host = MockHost()
        self.assertIsNotNone(host)
        self.assertIsNone(host._servo_state)
        self.assertIsNotNone(host.get_servo_state())
        self.assertEqual(host.get_servo_state(), servo_host.SERVO_STATE_BROKEN)

    def test_verify_set_state_broken_if_raised_error(self):
        host = MockHost()
        host._is_localhost = True
        host._repair_strategy = mock.Mock()
        host._repair_strategy.verify.side_effect = Exception('something_ex')
        try:
            host.verify(silent=True)
            self.assertEqual("Should not be reached", 'expecting error')
        except:
            pass
        self.assertEqual(host.get_servo_state(), servo_host.SERVO_STATE_BROKEN)

    def test_verify_set_state_working_if_no_raised_error(self):
        host = MockHost()
        host._repair_strategy = mock.Mock()
        host.verify(silent=True)
        self.assertEqual(host.get_servo_state(), servo_host.SERVO_STATE_WORKING)

    def test_repair_set_state_broken_if_raised_error(self):
        host = MockHost()
        host._is_localhost = True
        host._repair_strategy = mock.Mock()
        host._repair_strategy.repair.side_effect = Exception('something_ex')
        try:
            host.repair(silent=True)
            self.assertEqual("Should not be reached", 'expecting error')
        except:
            pass
        self.assertEqual(host.get_servo_state(), servo_host.SERVO_STATE_BROKEN)

    def test_repair_set_state_working_if_no_raised_error(self):
        host = MockHost()
        host._is_labstation = False
        host._repair_strategy = mock.Mock()
        host.repair(silent=True)
        self.assertEqual(host.get_servo_state(), servo_host.SERVO_STATE_WORKING)


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