1#!/usr/bin/python
2#
3# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
4# Use of this source code is governed by a BSD-style license that can be
5# found in the LICENSE file.
6
7"""Unit tests for server/cros/dynamic_suite/tools.py."""
8
9import mox
10import unittest
11
12import common
13from autotest_lib.server.cros.dynamic_suite.fakes import FakeHost
14from autotest_lib.server.cros.dynamic_suite.host_spec import HostSpec
15from autotest_lib.server.cros.dynamic_suite import host_spec
16from autotest_lib.server.cros.dynamic_suite import tools
17from autotest_lib.server import frontend
18
19
20class DynamicSuiteToolsTest(mox.MoxTestBase):
21    """Unit tests for dynamic_suite tools module methods.
22
23    @var _BOARD: fake board to reimage
24    """
25
26    _BOARD = 'board'
27    _DEPENDENCIES = {'test1': ['label1'], 'test2': ['label2']}
28    _POOL = 'bvt'
29
30    def setUp(self):
31        super(DynamicSuiteToolsTest, self).setUp()
32        self.afe = self.mox.CreateMock(frontend.AFE)
33        self.tko = self.mox.CreateMock(frontend.TKO)
34        # Having these ordered by complexity is important!
35        host_spec_list = [HostSpec([self._BOARD, self._POOL])]
36        for dep_list in self._DEPENDENCIES.itervalues():
37            host_spec_list.append(
38                HostSpec([self._BOARD, self._POOL], dep_list))
39        self.specs = host_spec.order_by_complexity(host_spec_list)
40
41    def testInjectVars(self):
42        """Should inject dict of varibles into provided strings."""
43        def find_all_in(d, s):
44            """Returns true if all key-value pairs in |d| are printed in |s|
45            and the dictionary representation is also in |s|.
46
47            @param d: the variable dictionary to check.
48            @param s: the control file string.
49            """
50            for k, v in d.iteritems():
51                if isinstance(v, str):
52                    if "%s='%s'\n" % (k, v) not in s:
53                        return False
54                else:
55                    if "%s=%r\n" % (k, v) not in s:
56                        return False
57            args_dict_str = "%s=%s\n" % ('args_dict', repr(d))
58            if args_dict_str not in s:
59                return False
60            return True
61
62        v = {'v1': 'one', 'v2': 'two', 'v3': None, 'v4': False, 'v5': 5}
63        self.assertTrue(find_all_in(v, tools.inject_vars(v, '')))
64        self.assertTrue(find_all_in(v, tools.inject_vars(v, 'ctrl')))
65        control_file = tools.inject_vars(v, 'sample')
66        self.assertTrue(tools._INJECT_BEGIN in control_file)
67        self.assertTrue(tools._INJECT_END in control_file)
68
69    def testRemoveInjection(self):
70        """Tests remove the injected variables from control file."""
71        control_file = """
72# INJECT_BEGIN - DO NOT DELETE THIS LINE
73v1='one'
74v4=False
75v5=5
76args_dict={'v1': 'one', 'v2': 'two', 'v3': None, 'v4': False, 'v5': 5}
77# INJECT_END - DO NOT DELETE LINE
78def init():
79    pass
80        """
81        control_file = tools.remove_injection(control_file)
82        self.assertTrue(control_file.strip().startswith('def init():'))
83
84    def testRemoveLegacyInjection(self):
85        """Tests remove the injected variables from legacy control file."""
86        control_file = """
87v1='one'
88_v2=False
89v3_x11_=5
90args_dict={'v1': 'one', '_v2': False, 'v3_x11': 5}
91def init():
92    pass
93        """
94        control_file = tools.remove_legacy_injection(control_file)
95        self.assertTrue(control_file.strip().startswith('def init():'))
96        control_file = tools.remove_injection(control_file)
97        self.assertTrue(control_file.strip().startswith('def init():'))
98
99
100    def testIncorrectlyLocked(self):
101        """Should detect hosts locked by random users."""
102        host = FakeHost(locked=True, locked_by='some guy')
103        self.assertTrue(tools.incorrectly_locked(host))
104
105
106    def testNotIncorrectlyLocked(self):
107        """Should accept hosts locked by the infrastructure."""
108        infra_user = 'an infra user'
109        self.mox.StubOutWithMock(tools, 'infrastructure_user')
110        tools.infrastructure_user().AndReturn(infra_user)
111        self.mox.ReplayAll()
112        host = FakeHost(locked=True, locked_by=infra_user)
113        self.assertFalse(tools.incorrectly_locked(host))
114
115
116if __name__ == "__main__":
117    unittest.main()
118