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 6 import unittest 7 8 import common 9 from autotest_lib.server import afe_utils 10 from autotest_lib.server import site_utils 11 12 13 class 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 26 class 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 61 if __name__ == '__main__': 62 unittest.main() 63 64
