• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/python
2# Copyright 2016 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6import mock
7import unittest
8
9import common
10
11from autotest_lib.server import utils
12from autotest_lib.server.hosts import host_info
13from autotest_lib.server.hosts import base_label
14
15 # pylint: disable=missing-docstring
16
17
18class TestBaseLabel(base_label.BaseLabel):
19    """TestBaseLabel is used for testing/validating BaseLabel methods."""
20
21    _NAME = 'base_label'
22
23    def exists(self, host):
24        return host.exists
25
26
27class TestBaseLabels(base_label.BaseLabel):
28    """
29    TestBaseLabels is used for testing/validating BaseLabel methods.
30
31    This is a variation of BaseLabel with multiple labels for _NAME
32    to ensure we handle a label that contains a list of labels for
33    its _NAME attribute.
34    """
35
36    _NAME = ['base_label_1' , 'base_label_2']
37
38
39class TestStringPrefixLabel(base_label.StringPrefixLabel):
40    """
41    TestBaseLabels is used for testing/validating StringPrefixLabel methods.
42
43    This test class is to check that we properly construct the prefix labels
44    with the label passed in during construction.
45    """
46
47    _NAME = 'prefix'
48
49    def __init__(self, label='postfix'):
50        self.label_to_return = label
51
52
53    def generate_labels(self, _):
54        return [self.label_to_return]
55
56
57class MockAFEHost(utils.EmptyAFEHost):
58
59    def __init__(self, labels=None, attributes=None):
60        self.labels = labels or []
61        self.attributes = attributes or {}
62
63
64class MockHost(object):
65
66    def __init__(self, exists=True, store=None):
67        self.hostname = 'hostname'
68        self.exists = exists
69        self.host_info_store = store
70
71
72class BaseLabelUnittests(unittest.TestCase):
73    """Unittest for testing base_label.BaseLabel."""
74
75    def setUp(self):
76        self.test_base_label = TestBaseLabel()
77        self.test_base_labels = TestBaseLabels()
78
79
80    def test_generate_labels(self):
81        """Let's make sure generate_labels() returns the labels expected."""
82        self.assertEqual(self.test_base_label.generate_labels(None),
83                         [self.test_base_label._NAME])
84
85
86    def test_get(self):
87        """Let's make sure the logic in get() works as expected."""
88        # We should get labels here.
89        self.assertEqual(self.test_base_label.get(MockHost(exists=True)),
90                         [self.test_base_label._NAME])
91        # We should get nothing here.
92        self.assertEqual(self.test_base_label.get(MockHost(exists=False)),
93                         [])
94
95
96    def test_get_all_labels(self):
97        """Check that we get the expected labels for get_all_labels()."""
98        prefix_tbl, full_tbl = self.test_base_label.get_all_labels()
99        prefix_tbls, full_tbls = self.test_base_labels.get_all_labels()
100
101        # We want to check that we always get a list of labels regardless if
102        # the label class attribute _NAME is a list or a string.
103        self.assertEqual(full_tbl, set([self.test_base_label._NAME]))
104        self.assertEqual(full_tbls, set(self.test_base_labels._NAME))
105
106        # We want to make sure we get nothing on the prefix_* side of things
107        # since BaseLabel shouldn't be a prefix for any label.
108        self.assertEqual(prefix_tbl, set())
109        self.assertEqual(prefix_tbls, set())
110
111
112class StringPrefixLabelUnittests(unittest.TestCase):
113    """Unittest for testing base_label.StringPrefixLabel."""
114
115    def setUp(self):
116        self.postfix_label = 'postfix_label'
117        self.test_label = TestStringPrefixLabel(label=self.postfix_label)
118
119
120    def test_get(self):
121        """Let's make sure that the labels we get are prefixed."""
122        self.assertEqual(self.test_label.get(None),
123                         ['%s:%s' % (self.test_label._NAME,
124                                     self.postfix_label)])
125
126
127    def test_get_all_labels(self):
128        """Check that we only get prefix labels and no full labels."""
129        prefix_labels, postfix_labels = self.test_label.get_all_labels()
130        self.assertEqual(prefix_labels, set(['%s:' % self.test_label._NAME]))
131        self.assertEqual(postfix_labels, set())
132
133
134class LabelRetrieverUnittests(unittest.TestCase):
135    """Unittest for testing base_label.LabelRetriever."""
136
137    def setUp(self):
138        label_list = [TestStringPrefixLabel(), TestBaseLabel()]
139        self.retriever = base_label.LabelRetriever(label_list)
140        self.retriever._populate_known_labels(label_list)
141
142
143    def test_populate_known_labels(self):
144        """Check that _populate_known_labels() works as expected."""
145        full_names = set([TestBaseLabel._NAME])
146        prefix_names = set(['%s:' % TestStringPrefixLabel._NAME])
147        # Check on a normal retriever.
148        self.assertEqual(self.retriever.label_full_names, full_names)
149        self.assertEqual(self.retriever.label_prefix_names, prefix_names)
150
151
152    def test_is_known_label(self):
153        """Check _is_known_label() detects/skips the right labels."""
154        # This will be a list of tuples of label and expected return bool.
155        # Make sure Full matches match correctly
156        labels_to_check = [(TestBaseLabel._NAME, True),
157                           ('%s:' % TestStringPrefixLabel._NAME, True),
158                           # Make sure partial matches fail.
159                           (TestBaseLabel._NAME[:2], False),
160                           ('%s:' % TestStringPrefixLabel._NAME[:2], False),
161                           ('no_label_match', False)]
162
163        for label, expected_known in labels_to_check:
164            self.assertEqual(self.retriever._is_known_label(label),
165                             expected_known)
166
167
168    def test_update_labels(self):
169        """Check that we add/remove the expected labels in update_labels()."""
170        label_to_add = 'label_to_add'
171        label_to_remove = 'prefix:label_to_remove'
172        store = host_info.InMemoryHostInfoStore(
173                info=host_info.HostInfo(
174                        labels=[label_to_remove, TestBaseLabel._NAME],
175                ),
176        )
177        mockhost = MockHost(store=store)
178
179        retriever = base_label.LabelRetriever(
180                [TestStringPrefixLabel(label=label_to_add),
181                 TestBaseLabel()])
182        retriever.update_labels(mockhost)
183        self.assertEqual(
184                set(store.get().labels),
185                {'%s:%s' % (TestStringPrefixLabel._NAME, label_to_add),
186                 TestBaseLabel._NAME},
187        )
188
189
190if __name__ == '__main__':
191    unittest.main()
192
193