1# Copyright 2020 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5
6from __future__ import print_function
7from __future__ import absolute_import
8from __future__ import division
9from __future__ import unicode_literals
10
11import six
12import unittest
13
14import common
15import autotest_lib.site_utils.deployment.prepare.dut as dut
16
17from autotest_lib.server.hosts.host_info import InMemoryHostInfoStore, HostInfo
18
19
20mock_labstation_args = {
21    "labels": [
22        "os:labstation",
23    ],
24    "attributes": None,
25    "stable_versions": None,
26}
27
28mock_labstation_args_bad_config = {
29    "labels": [
30        "os:labstation",
31    ],
32    "attributes": {
33        "powerunit_outlet": "AAAA"
34    },
35    "stable_versions": None,
36}
37
38
39class MockHost(object):
40    """Mock Host object, mimics a labstation with a servo attached"""
41    SHUTDOWN_TIMEOUT = 0.001
42    BOOT_TIMEOUT = 0.001
43
44    def __init__(self, info, power_on=True, power_off=True):
45        self.host_info_store = InMemoryHostInfoStore(info=info)
46        self.hostname = "a"
47        self._called_rpm_power_off_and_wait = 0
48        self._called_rpm_power_on_and_wait = 0
49
50    def rpm_power_off_and_wait(self):
51        self._called_rpm_power_off_and_wait += 1
52
53    def rpm_power_on_and_wait(self):
54        self._called_rpm_power_on_and_wait += 1
55
56class TestDut(unittest.TestCase):
57    def test_verify_labstation_RPM_config_unsafe(self):
58        """Test successful labstation RPM powercycle"""
59        m = MockHost(info=HostInfo(**mock_labstation_args))
60        dut.verify_labstation_RPM_config_unsafe(m)
61
62    def test_verify_labstation_RPM_config_unsafe_bad_config(self):
63        m = MockHost(info=HostInfo(**mock_labstation_args_bad_config))
64        with self.assertRaises(Exception):
65            dut.verify_labstation_RPM_config_unsafe(m)
66
67
68if __name__ == "__main__":
69    unittest.main()
70