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 unittest
7
8import common
9from autotest_lib.server import afe_utils
10from autotest_lib.server import site_utils
11
12
13class MockHost(object):
14    """
15    Object to represent host used to test afe_util.py methods.
16    """
17
18    def __init__(self, labels=[]):
19      """
20      Setup the self._afe_host attribute since that's what we're mostly using.
21      """
22      self._afe_host = site_utils.EmptyAFEHost()
23      self._afe_host.labels = labels
24
25
26class AfeUtilsUnittest(unittest.TestCase):
27    """
28    Test functions in afe_utils.py.
29    """
30
31    def testGetLabels(self):
32        """
33        Test method get_labels returns expected labels.
34        """
35        prefix = 'prefix'
36        expected_labels = [prefix + ':' + str(i) for i in range(5)]
37        all_labels = []
38        all_labels += expected_labels
39        all_labels += [str(i) for i in range(6, 9)]
40        host = MockHost(labels=all_labels)
41        got_labels = afe_utils.get_labels(host, prefix)
42
43        self.assertItemsEqual(got_labels, expected_labels)
44
45
46    def testGetLabelsAll(self):
47        """
48        Test method get_labels returns all labels.
49        """
50        prefix = 'prefix'
51        prefix_labels = [prefix + ':' + str(i) for i in range(5)]
52        all_labels = []
53        all_labels += prefix_labels
54        all_labels += [str(i) for i in range(6, 9)]
55        host = MockHost(labels=all_labels)
56        got_labels = afe_utils.get_labels(host)
57
58        self.assertItemsEqual(got_labels, all_labels)
59
60
61if __name__ == '__main__':
62    unittest.main()
63
64