blob: 67b08942697797e81e56800fe72a16193bfedbd9 [file] [log] [blame]
# Copyright 2016 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
import mock
import unittest
import common
from autotest_lib.server.cros.dynamic_suite import frontend_wrappers
from autotest_lib.server import utils
from autotest_lib.server.hosts import base_label
# pylint: disable=missing-docstring
class TestBaseLabel(base_label.BaseLabel):
"""TestBaseLabel is used for testing/validating BaseLabel methods."""
_NAME = 'base_label'
def exists(self, host):
return host.exists
class TestBaseLabels(base_label.BaseLabel):
TestBaseLabels is used for testing/validating BaseLabel methods.
This is a variation of BaseLabel with multiple labels for _NAME
to ensure we handle a label that contains a list of labels for
its _NAME attribute.
_NAME = ['base_label_1' , 'base_label_2']
class TestStringPrefixLabel(base_label.StringPrefixLabel):
TestBaseLabels is used for testing/validating StringPrefixLabel methods.
This test class is to check that we properly construct the prefix labels
with the label passed in during construction.
_NAME = 'prefix'
def __init__(self, label='postfix'):
self.label_to_return = label
def generate_labels(self, _):
return [self.label_to_return]
class TestLabelList(base_label.StringLabel):
"""TestLabelList is used to validate a label consisting of other labels."""
_LABEL_LIST = [TestBaseLabel(), TestStringPrefixLabel()]
def generate_labels(self, host):
labels = []
for label_detectors in self._LABEL_LIST:
return labels
class MockAFEHost(utils.EmptyAFEHost):
def __init__(self, labels=[], attributes={}):
self.labels = labels
self.attributes = attributes
class MockHost(object):
def __init__(self, exists=True, afe_host=None):
self.hostname = 'hostname'
self.exists = exists
self._afe_host = afe_host
class BaseLabelUnittests(unittest.TestCase):
"""Unittest for testing base_label.BaseLabel."""
def setUp(self):
self.test_base_label = TestBaseLabel()
self.test_base_labels = TestBaseLabels()
def test_generate_labels(self):
"""Let's make sure generate_labels() returns the labels expected."""
def test_get(self):
"""Let's make sure the logic in get() works as expected."""
# We should get labels here.
# We should get nothing here.
def test_get_all_labels(self):
"""Check that we get the expected labels for get_all_labels()."""
prefix_tbl, full_tbl = self.test_base_label.get_all_labels()
prefix_tbls, full_tbls = self.test_base_labels.get_all_labels()
# We want to check that we always get a list of labels regardless if
# the label class attribute _NAME is a list or a string.
self.assertEqual(full_tbl, set([self.test_base_label._NAME]))
self.assertEqual(full_tbls, set(self.test_base_labels._NAME))
# We want to make sure we get nothing on the prefix_* side of things
# since BaseLabel shouldn't be a prefix for any label.
self.assertEqual(prefix_tbl, set())
self.assertEqual(prefix_tbls, set())
class StringPrefixLabelUnittests(unittest.TestCase):
"""Unittest for testing base_label.StringPrefixLabel."""
def setUp(self):
self.postfix_label = 'postfix_label'
self.test_label = TestStringPrefixLabel(label=self.postfix_label)
def test_get(self):
"""Let's make sure that the labels we get are prefixed."""
['%s:%s' % (self.test_label._NAME,
def test_get_all_labels(self):
"""Check that we only get prefix labels and no full labels."""
prefix_labels, postfix_labels = self.test_label.get_all_labels()
self.assertEqual(prefix_labels, set(['%s:' % self.test_label._NAME]))
self.assertEqual(postfix_labels, set())
class LabelRetrieverUnittests(unittest.TestCase):
"""Unittest for testing base_label.LabelRetriever."""
def setUp(self):
label_list = [TestStringPrefixLabel(), TestBaseLabel()]
self.retriever = base_label.LabelRetriever(label_list)
self.retriever_label_list = base_label.LabelRetriever([TestLabelList()])
def test_populate_known_labels(self):
"""Check that _populate_known_labels() works as expected."""
full_names = set([TestBaseLabel._NAME])
prefix_names = set(['%s:' % TestStringPrefixLabel._NAME])
# Check on a normal retriever.
self.assertEqual(self.retriever.label_full_names, full_names)
self.assertEqual(self.retriever.label_prefix_names, prefix_names)
# Check on a retriever that has a label with a label list.
self.assertEqual(self.retriever_label_list.label_full_names, full_names)
def test_is_known_label(self):
"""Check _is_known_label() detects/skips the right labels."""
# This will be a list of tuples of label and expected return bool.
# Make sure Full matches match correctly
labels_to_check = [(TestBaseLabel._NAME, True),
('%s:' % TestStringPrefixLabel._NAME, True),
# Make sure partial matches fail.
(TestBaseLabel._NAME[:2], False),
('%s:' % TestStringPrefixLabel._NAME[:2], False),
('no_label_match', False)]
for label, expected_known in labels_to_check:
@mock.patch.object(frontend_wrappers, 'RetryingAFE')
def test_update_labels(self, mock_retry_afe):
"""Check that we add/remove the expected labels in update_labels()."""
label_to_add = 'label_to_add'
label_to_remove = 'prefix:label_to_remove'
mock_afe = mock.MagicMock()
mockhost = MockHost(afe_host=MockAFEHost(
labels=[label_to_remove, TestBaseLabel._NAME]))
expected_remove_labels = [label_to_remove]
expected_add_labels = ['%s:%s' % (TestStringPrefixLabel._NAME,
retriever = base_label.LabelRetriever(
# Check that we removed the right labels'host_remove_labels',
# Check that we added the right labels'host_add_labels',
if __name__ == '__main__':