Home | History | Annotate | Download | only in site_utils
      1 #!/usr/bin/env python
      2 # Copyright 2015 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 collections
      7 import itertools
      8 import logging
      9 import os
     10 import unittest
     11 
     12 import common
     13 from autotest_lib.server.lib import status_history
     14 from autotest_lib.site_utils import lab_inventory
     15 
     16 
     17 class _FakeHost(object):
     18     """Class to mock `Host` in _FakeHostHistory for testing."""
     19 
     20     def __init__(self, hostname):
     21         self.hostname = hostname
     22 
     23 
     24 class _FakeHostHistory(object):
     25     """Class to mock `HostJobHistory` for testing."""
     26 
     27     def __init__(self, model, pool, status, hostname=''):
     28         self.host_model = model
     29         self.host_board = model + '_board'
     30         self.host_pool = pool
     31         self.status = status
     32         self.host = _FakeHost(hostname)
     33 
     34 
     35     def last_diagnosis(self):
     36         """Return the recorded diagnosis."""
     37         return self.status, None
     38 
     39 
     40 class _FakeHostLocation(object):
     41     """Class to mock `HostJobHistory` for location sorting."""
     42 
     43     _HOSTNAME_FORMAT = 'chromeos%d-row%d-rack%d-host%d'
     44 
     45 
     46     def __init__(self, location):
     47         self.hostname = self._HOSTNAME_FORMAT % location
     48 
     49 
     50     @property
     51     def host(self):
     52         """Return a fake host object with a hostname."""
     53         return self
     54 
     55 
     56 # Status values that may be returned by `HostJobHistory`.
     57 #
     58 # These merely rename the corresponding values in `status_history`
     59 # for convenience.
     60 
     61 _WORKING = status_history.WORKING
     62 _UNUSED = status_history.UNUSED
     63 _BROKEN = status_history.BROKEN
     64 _UNKNOWN = status_history.UNKNOWN
     65 
     66 
     67 class HostSetInventoryTestCase(unittest.TestCase):
     68     """Unit tests for class `_HostSetInventory`.
     69 
     70     Coverage is quite basic:  mostly just enough to make sure every
     71     function gets called, and to make sure that the counting knows
     72     the difference between 0 and 1.
     73 
     74     The testing also ensures that all known status values that can be
     75     returned by `HostJobHistory` are counted as expected.
     76     """
     77 
     78     def setUp(self):
     79         super(HostSetInventoryTestCase, self).setUp()
     80         self.histories = lab_inventory._HostSetInventory()
     81 
     82 
     83     def _add_host(self, status):
     84         fake = _FakeHostHistory('zebra', lab_inventory.SPARE_POOL, status)
     85         self.histories.record_host(fake)
     86 
     87 
     88     def _check_counts(self, working, broken, idle):
     89         """Check that pool counts match expectations.
     90 
     91         Asserts that `get_working()`, `get_broken()`, and `get_idle()`
     92         return the given expected values.  Also assert that
     93         `get_total()` is the sum of all counts.
     94 
     95         @param working The expected total of working devices.
     96         @param broken  The expected total of broken devices.
     97         @param idle  The expected total of idle devices.
     98         """
     99         self.assertEqual(self.histories.get_working(), working)
    100         self.assertEqual(self.histories.get_broken(), broken)
    101         self.assertEqual(self.histories.get_idle(), idle)
    102         self.assertEqual(self.histories.get_total(),
    103                          working + broken + idle)
    104 
    105 
    106     def test_empty(self):
    107         """Test counts when there are no DUTs recorded."""
    108         self._check_counts(0, 0, 0)
    109 
    110 
    111     def test_broken(self):
    112         """Test counting for broken DUTs."""
    113         self._add_host(_BROKEN)
    114         self._check_counts(0, 1, 0)
    115 
    116 
    117     def test_working(self):
    118         """Test counting for working DUTs."""
    119         self._add_host(_WORKING)
    120         self._check_counts(1, 0, 0)
    121 
    122 
    123     def test_idle(self):
    124         """Testing counting for idle status values."""
    125         self._add_host(_UNUSED)
    126         self._check_counts(0, 0, 1)
    127         self._add_host(_UNKNOWN)
    128         self._check_counts(0, 0, 2)
    129 
    130 
    131     def test_working_then_broken(self):
    132         """Test counts after adding a working and then a broken DUT."""
    133         self._add_host(_WORKING)
    134         self._add_host(_BROKEN)
    135         self._check_counts(1, 1, 0)
    136 
    137 
    138     def test_broken_then_working(self):
    139         """Test counts after adding a broken and then a working DUT."""
    140         self._add_host(_BROKEN)
    141         self._add_host(_WORKING)
    142         self._check_counts(1, 1, 0)
    143 
    144 
    145 class PoolSetInventoryTestCase(unittest.TestCase):
    146     """Unit tests for class `_PoolSetInventory`.
    147 
    148     Coverage is quite basic:  just enough to make sure every function
    149     gets called, and to make sure that the counting knows the difference
    150     between 0 and 1.
    151 
    152     The tests make sure that both individual pool counts and totals are
    153     counted correctly.
    154     """
    155 
    156     _POOL_SET = ['humpty', 'dumpty']
    157 
    158     def setUp(self):
    159         super(PoolSetInventoryTestCase, self).setUp()
    160         self._pool_histories = lab_inventory._PoolSetInventory(self._POOL_SET)
    161 
    162 
    163     def _add_host(self, pool, status):
    164         fake = _FakeHostHistory('zebra', pool, status)
    165         self._pool_histories.record_host(fake)
    166 
    167 
    168     def _check_all_counts(self, working, broken):
    169         """Check that total counts for all pools match expectations.
    170 
    171         Checks that `get_working()` and `get_broken()` return the
    172         given expected values when called without a pool specified.
    173         Also check that `get_total()` is the sum of working and
    174         broken devices.
    175 
    176         Additionally, call the various functions for all the pools
    177         individually, and confirm that the totals across pools match
    178         the given expectations.
    179 
    180         @param working The expected total of working devices.
    181         @param broken  The expected total of broken devices.
    182 
    183         """
    184         self.assertEqual(self._pool_histories.get_working(), working)
    185         self.assertEqual(self._pool_histories.get_broken(), broken)
    186         self.assertEqual(self._pool_histories.get_total(),
    187                          working + broken)
    188         count_working = 0
    189         count_broken = 0
    190         count_total = 0
    191         for pool in self._POOL_SET:
    192             count_working += self._pool_histories.get_working(pool)
    193             count_broken += self._pool_histories.get_broken(pool)
    194             count_total += self._pool_histories.get_total(pool)
    195         self.assertEqual(count_working, working)
    196         self.assertEqual(count_broken, broken)
    197         self.assertEqual(count_total, working + broken)
    198 
    199 
    200     def _check_pool_counts(self, pool, working, broken):
    201         """Check that counts for a given pool match expectations.
    202 
    203         Checks that `get_working()` and `get_broken()` return the
    204         given expected values for the given pool.  Also check that
    205         `get_total()` is the sum of working and broken devices.
    206 
    207         @param pool    The pool to be checked.
    208         @param working The expected total of working devices.
    209         @param broken  The expected total of broken devices.
    210 
    211         """
    212         self.assertEqual(self._pool_histories.get_working(pool),
    213                          working)
    214         self.assertEqual(self._pool_histories.get_broken(pool),
    215                          broken)
    216         self.assertEqual(self._pool_histories.get_total(pool),
    217                          working + broken)
    218 
    219 
    220     def test_empty(self):
    221         """Test counts when there are no DUTs recorded."""
    222         self._check_all_counts(0, 0)
    223         for pool in self._POOL_SET:
    224             self._check_pool_counts(pool, 0, 0)
    225 
    226 
    227     def test_all_working_then_broken(self):
    228         """Test counts after adding a working and then a broken DUT.
    229 
    230         For each pool, add first a working, then a broken DUT.  After
    231         each DUT is added, check counts to confirm the correct values.
    232         """
    233         working = 0
    234         broken = 0
    235         for pool in self._POOL_SET:
    236             self._add_host(pool, _WORKING)
    237             working += 1
    238             self._check_pool_counts(pool, 1, 0)
    239             self._check_all_counts(working, broken)
    240             self._add_host(pool, _BROKEN)
    241             broken += 1
    242             self._check_pool_counts(pool, 1, 1)
    243             self._check_all_counts(working, broken)
    244 
    245 
    246     def test_all_broken_then_working(self):
    247         """Test counts after adding a broken and then a working DUT.
    248 
    249         For each pool, add first a broken, then a working DUT.  After
    250         each DUT is added, check counts to confirm the correct values.
    251         """
    252         working = 0
    253         broken = 0
    254         for pool in self._POOL_SET:
    255             self._add_host(pool, _BROKEN)
    256             broken += 1
    257             self._check_pool_counts(pool, 0, 1)
    258             self._check_all_counts(working, broken)
    259             self._add_host(pool, _WORKING)
    260             working += 1
    261             self._check_pool_counts(pool, 1, 1)
    262             self._check_all_counts(working, broken)
    263 
    264 
    265 class LocationSortTests(unittest.TestCase):
    266     """Unit tests for `_sort_by_location()`."""
    267 
    268     def setUp(self):
    269         super(LocationSortTests, self).setUp()
    270 
    271 
    272     def _check_sorting(self, *locations):
    273         """Test sorting a given list of locations.
    274 
    275         The input is an already ordered list of lists of tuples with
    276         row, rack, and host numbers.  The test converts the tuples
    277         to hostnames, preserving the original ordering.  Then it
    278         flattens and scrambles the input, runs it through
    279         `_sort_by_location()`, and asserts that the result matches
    280         the original.
    281 
    282         """
    283         lab = 0
    284         expected = []
    285         for tuples in locations:
    286             lab += 1
    287             expected.append(
    288                     [_FakeHostLocation((lab,) + t) for t in tuples])
    289         scrambled = [e for e in itertools.chain(*expected)]
    290         scrambled = [e for e in reversed(scrambled)]
    291         actual = lab_inventory._sort_by_location(scrambled)
    292         # The ordering of the labs in the output isn't guaranteed,
    293         # so we can't compare `expected` and `actual` directly.
    294         # Instead, we create a dictionary keyed on the first host in
    295         # each lab, and compare the dictionaries.
    296         self.assertEqual({l[0]: l for l in expected},
    297                          {l[0]: l for l in actual})
    298 
    299 
    300     def test_separate_labs(self):
    301         """Test that sorting distinguishes labs."""
    302         self._check_sorting([(1, 1, 1)], [(1, 1, 1)], [(1, 1, 1)])
    303 
    304 
    305     def test_separate_rows(self):
    306         """Test for proper sorting when only rows are different."""
    307         self._check_sorting([(1, 1, 1), (9, 1, 1), (10, 1, 1)])
    308 
    309 
    310     def test_separate_racks(self):
    311         """Test for proper sorting when only racks are different."""
    312         self._check_sorting([(1, 1, 1), (1, 9, 1), (1, 10, 1)])
    313 
    314 
    315     def test_separate_hosts(self):
    316         """Test for proper sorting when only hosts are different."""
    317         self._check_sorting([(1, 1, 1), (1, 1, 9), (1, 1, 10)])
    318 
    319 
    320     def test_diagonal(self):
    321         """Test for proper sorting when all parts are different."""
    322         self._check_sorting([(1, 1, 2), (1, 2, 1), (2, 1, 1)])
    323 
    324 
    325 class InventoryScoringTests(unittest.TestCase):
    326     """Unit tests for `_score_repair_set()`."""
    327 
    328     def setUp(self):
    329         super(InventoryScoringTests, self).setUp()
    330 
    331 
    332     def _make_buffer_counts(self, *counts):
    333         """Create a dictionary suitable as `buffer_counts`.
    334 
    335         @param counts List of tuples with model count data.
    336         """
    337         self._buffer_counts = dict(counts)
    338 
    339 
    340     def _make_history_list(self, repair_counts):
    341         """Create a list suitable as `repair_list`.
    342 
    343         @param repair_counts List of (model, count) tuples.
    344         """
    345         pool = lab_inventory.SPARE_POOL
    346         histories = []
    347         for model, count in repair_counts:
    348             for i in range(0, count):
    349                 histories.append(
    350                     _FakeHostHistory(model, pool, _BROKEN))
    351         return histories
    352 
    353 
    354     def _check_better(self, repair_a, repair_b):
    355         """Test that repair set A scores better than B.
    356 
    357         Contruct repair sets from `repair_a` and `repair_b`,
    358         and score both of them using the pre-existing
    359         `self._buffer_counts`.  Assert that the score for A is
    360         better than the score for B.
    361 
    362         @param repair_a Input data for repair set A
    363         @param repair_b Input data for repair set B
    364         """
    365         score_a = lab_inventory._score_repair_set(
    366                 self._buffer_counts,
    367                 self._make_history_list(repair_a))
    368         score_b = lab_inventory._score_repair_set(
    369                 self._buffer_counts,
    370                 self._make_history_list(repair_b))
    371         self.assertGreater(score_a, score_b)
    372 
    373 
    374     def _check_equal(self, repair_a, repair_b):
    375         """Test that repair set A scores the same as B.
    376 
    377         Contruct repair sets from `repair_a` and `repair_b`,
    378         and score both of them using the pre-existing
    379         `self._buffer_counts`.  Assert that the score for A is
    380         equal to the score for B.
    381 
    382         @param repair_a Input data for repair set A
    383         @param repair_b Input data for repair set B
    384         """
    385         score_a = lab_inventory._score_repair_set(
    386                 self._buffer_counts,
    387                 self._make_history_list(repair_a))
    388         score_b = lab_inventory._score_repair_set(
    389                 self._buffer_counts,
    390                 self._make_history_list(repair_b))
    391         self.assertEqual(score_a, score_b)
    392 
    393 
    394     def test_improve_worst_model(self):
    395         """Test that improving the worst model improves scoring.
    396 
    397         Construct a buffer counts dictionary with all models having
    398         different counts.  Assert that it is both necessary and
    399         sufficient to improve the count of the worst model in order
    400         to improve the score.
    401         """
    402         self._make_buffer_counts(('lion', 0),
    403                                  ('tiger', 1),
    404                                  ('bear', 2))
    405         self._check_better([('lion', 1)], [('tiger', 1)])
    406         self._check_better([('lion', 1)], [('bear', 1)])
    407         self._check_better([('lion', 1)], [('tiger', 2)])
    408         self._check_better([('lion', 1)], [('bear', 2)])
    409         self._check_equal([('tiger', 1)], [('bear', 1)])
    410 
    411 
    412     def test_improve_worst_case_count(self):
    413         """Test that improving the number of worst cases improves the score.
    414 
    415         Construct a buffer counts dictionary with all models having
    416         the same counts.  Assert that improving two models is better
    417         than improving one.  Assert that improving any one model is
    418         as good as any other.
    419         """
    420         self._make_buffer_counts(('lion', 0),
    421                                  ('tiger', 0),
    422                                  ('bear', 0))
    423         self._check_better([('lion', 1), ('tiger', 1)], [('bear', 2)])
    424         self._check_equal([('lion', 2)], [('tiger', 1)])
    425         self._check_equal([('tiger', 1)], [('bear', 1)])
    426 
    427 
    428 # Each item is the number of DUTs in that status.
    429 STATUS_CHOICES = (_WORKING, _BROKEN, _UNUSED)
    430 StatusCounts = collections.namedtuple('StatusCounts', ['good', 'bad', 'idle'])
    431 # Each item is a StatusCounts tuple specifying the number of DUTs per status in
    432 # the that pool.
    433 CRITICAL_POOL = lab_inventory.CRITICAL_POOLS[0]
    434 SPARE_POOL = lab_inventory.SPARE_POOL
    435 POOL_CHOICES = (CRITICAL_POOL, SPARE_POOL)
    436 PoolStatusCounts = collections.namedtuple('PoolStatusCounts',
    437                                           ['critical', 'spare'])
    438 
    439 def create_inventory(data):
    440     """Create a `_LabInventory` instance for testing.
    441 
    442     This function allows the construction of a complete `_LabInventory`
    443     object from a simplified input representation.
    444 
    445     A single 'critical pool' is arbitrarily chosen for purposes of
    446     testing; there's no coverage for testing arbitrary combinations
    447     in more than one critical pool.
    448 
    449     @param data: dict {key: PoolStatusCounts}.
    450     @returns: lab_inventory._LabInventory object.
    451     """
    452     histories = []
    453     for model, counts in data.iteritems():
    454         for p, pool in enumerate(POOL_CHOICES):
    455             for s, status in enumerate(STATUS_CHOICES):
    456                 fake_host = _FakeHostHistory(model, pool, status)
    457                 histories.extend([fake_host] * counts[p][s])
    458     inventory = lab_inventory._LabInventory(
    459             histories, lab_inventory.MANAGED_POOLS)
    460     return inventory
    461 
    462 
    463 class LabInventoryTests(unittest.TestCase):
    464     """Tests for the basic functions of `_LabInventory`.
    465 
    466     Contains basic coverage to show that after an inventory is created
    467     and DUTs with known status are added, the inventory counts match the
    468     counts of the added DUTs.
    469     """
    470 
    471     _MODEL_LIST = ['lion', 'tiger', 'bear'] # Oh, my!
    472 
    473 
    474     def _check_inventory_counts(self, inventory, data, msg=None):
    475         """Check that all counts in the inventory match `data`.
    476 
    477         This asserts that the actual counts returned by the various
    478         accessor functions for `inventory` match the values expected for
    479         the given `data` that created the inventory.
    480 
    481         @param inventory: _LabInventory object to check.
    482         @param data Inventory data to check against. Same type as
    483                 `create_inventory`.
    484         """
    485         self.assertEqual(set(inventory.keys()), set(data.keys()))
    486         for model, histories in inventory.iteritems():
    487             expected_counts = data[model]
    488             actual_counts = PoolStatusCounts(
    489                     StatusCounts(
    490                             histories.get_working(CRITICAL_POOL),
    491                             histories.get_broken(CRITICAL_POOL),
    492                             histories.get_idle(CRITICAL_POOL),
    493                     ),
    494                     StatusCounts(
    495                             histories.get_working(SPARE_POOL),
    496                             histories.get_broken(SPARE_POOL),
    497                             histories.get_idle(SPARE_POOL),
    498                     ),
    499             )
    500             self.assertEqual(actual_counts, expected_counts, msg)
    501 
    502             self.assertEqual(len(histories.get_working_list()),
    503                              sum([p.good for p in expected_counts]),
    504                              msg)
    505             self.assertEqual(len(histories.get_broken_list()),
    506                              sum([p.bad for p in expected_counts]),
    507                              msg)
    508             self.assertEqual(len(histories.get_idle_list()),
    509                              sum([p.idle for p in expected_counts]),
    510                              msg)
    511 
    512 
    513     def test_empty(self):
    514         """Test counts when there are no DUTs recorded."""
    515         inventory = create_inventory({})
    516         self.assertEqual(inventory.get_num_duts(), 0)
    517         self.assertEqual(inventory.get_boards(), set())
    518         self._check_inventory_counts(inventory, {})
    519         self.assertEqual(inventory.get_num_models(), 0)
    520 
    521 
    522     def _check_model_count(self, model_count):
    523         """Parameterized test for testing a specific number of models."""
    524         msg = '[model: %d]' % (model_count,)
    525         models = self._MODEL_LIST[:model_count]
    526         data = {
    527                 m: PoolStatusCounts(
    528                         StatusCounts(1, 1, 1),
    529                         StatusCounts(1, 1, 1),
    530                 )
    531                 for m in models
    532         }
    533         inventory = create_inventory(data)
    534         self.assertEqual(inventory.get_num_duts(), 6 * model_count, msg)
    535         self.assertEqual(inventory.get_num_models(), model_count, msg)
    536         for pool in [CRITICAL_POOL, SPARE_POOL]:
    537             self.assertEqual(set(inventory.get_pool_models(pool)),
    538                              set(models))
    539         self._check_inventory_counts(inventory, data, msg=msg)
    540 
    541 
    542     def test_model_counts(self):
    543         """Test counts for various numbers of models."""
    544         self.longMessage = True
    545         for model_count in range(0, len(self._MODEL_LIST)):
    546             self._check_model_count(model_count)
    547 
    548 
    549     def _check_single_dut_counts(self, critical, spare):
    550         """Parmeterized test for single dut counts."""
    551         self.longMessage = True
    552         counts = PoolStatusCounts(critical, spare)
    553         model = self._MODEL_LIST[0]
    554         data = {model: counts}
    555         msg = '[data: %s]' % (data,)
    556         inventory = create_inventory(data)
    557         self.assertEqual(inventory.get_num_duts(), 1, msg)
    558         self.assertEqual(inventory.get_num_models(), 1, msg)
    559         self._check_inventory_counts(inventory, data, msg=msg)
    560 
    561 
    562     def test_single_dut_counts(self):
    563         """Test counts when there is a single DUT per board, and it is good."""
    564         status_100 = StatusCounts(1, 0, 0)
    565         status_010 = StatusCounts(0, 1, 0)
    566         status_001 = StatusCounts(0, 0, 1)
    567         status_null = StatusCounts(0, 0, 0)
    568         self._check_single_dut_counts(status_100, status_null)
    569         self._check_single_dut_counts(status_010, status_null)
    570         self._check_single_dut_counts(status_001, status_null)
    571         self._check_single_dut_counts(status_null, status_100)
    572         self._check_single_dut_counts(status_null, status_010)
    573         self._check_single_dut_counts(status_null, status_001)
    574 
    575 
    576 # MODEL_MESSAGE_TEMPLATE -
    577 # This is a sample of the output text produced by
    578 # _generate_model_inventory_message().  This string is parsed by the
    579 # tests below to construct a sample inventory that should produce
    580 # the output, and then the output is generated and checked against
    581 # this original sample.
    582 #
    583 # Constructing inventories from parsed sample text serves two
    584 # related purposes:
    585 #   - It provides a way to see what the output should look like
    586 #     without having to run the script.
    587 #   - It helps make sure that a human being will actually look at
    588 #     the output to see that it's basically readable.
    589 # This should also help prevent test bugs caused by writing tests
    590 # that simply parrot the original output generation code.
    591 
    592 _MODEL_MESSAGE_TEMPLATE = '''
    593 Model                  Avail   Bad  Idle  Good Spare Total
    594 lion                      -1    13     2    11    12    26
    595 tiger                     -1     5     2     9     4    16
    596 bear                       0     5     2    10     5    17
    597 platypus                   4     2     2    20     6    24
    598 aardvark                   7     2     2     6     9    10
    599 '''
    600 
    601 
    602 class PoolSetInventoryTests(unittest.TestCase):
    603     """Tests for `_generate_model_inventory_message()`.
    604 
    605     The tests create various test inventories designed to match the
    606     counts in `_MODEL_MESSAGE_TEMPLATE`, and asserts that the
    607     generated message text matches the original message text.
    608 
    609     Message text is represented as a list of strings, split on the
    610     `'\n'` separator.
    611     """
    612 
    613     def setUp(self):
    614         self.maxDiff = None
    615         lines = [x.strip() for x in _MODEL_MESSAGE_TEMPLATE.split('\n') if
    616                  x.strip()]
    617         self._header, self._model_lines = lines[0], lines[1:]
    618         self._model_data = []
    619         for l in self._model_lines:
    620             items = l.split()
    621             model = items[0]
    622             bad, idle, good, spare = [int(x) for x in items[2:-1]]
    623             self._model_data.append((model, (good, bad, idle, spare)))
    624 
    625 
    626 
    627     def _make_minimum_spares(self, counts):
    628         """Create a counts tuple with as few spare DUTs as possible."""
    629         good, bad, idle, spares = counts
    630         if spares > bad + idle:
    631             return PoolStatusCounts(
    632                     StatusCounts(good + bad +idle - spares, 0, 0),
    633                     StatusCounts(spares - bad - idle, bad, idle),
    634             )
    635         elif spares < bad:
    636             return PoolStatusCounts(
    637                     StatusCounts(good, bad - spares, idle),
    638                     StatusCounts(0, spares, 0),
    639             )
    640         else:
    641             return PoolStatusCounts(
    642                     StatusCounts(good, 0, idle + bad - spares),
    643                     StatusCounts(0, bad, spares - bad),
    644             )
    645 
    646 
    647     def _make_maximum_spares(self, counts):
    648         """Create a counts tuple with as many spare DUTs as possible."""
    649         good, bad, idle, spares = counts
    650         if good > spares:
    651             return PoolStatusCounts(
    652                     StatusCounts(good - spares, bad, idle),
    653                     StatusCounts(spares, 0, 0),
    654             )
    655         elif good + bad > spares:
    656             return PoolStatusCounts(
    657                     StatusCounts(0, good + bad - spares, idle),
    658                     StatusCounts(good, spares - good, 0),
    659             )
    660         else:
    661             return PoolStatusCounts(
    662                     StatusCounts(0, 0, good + bad + idle - spares),
    663                     StatusCounts(good, bad, spares - good - bad),
    664             )
    665 
    666 
    667     def _check_message(self, message):
    668         """Checks that message approximately matches expected string."""
    669         message = [x.strip() for x in message.split('\n') if x.strip()]
    670         self.assertIn(self._header, message)
    671         body = message[message.index(self._header) + 1:]
    672         self.assertEqual(body, self._model_lines)
    673 
    674 
    675     def test_minimum_spares(self):
    676         """Test message generation when the spares pool is low."""
    677         data = {
    678             model: self._make_minimum_spares(counts)
    679                 for model, counts in self._model_data
    680         }
    681         inventory = create_inventory(data)
    682         message = lab_inventory._generate_model_inventory_message(inventory)
    683         self._check_message(message)
    684 
    685     def test_maximum_spares(self):
    686         """Test message generation when the critical pool is low."""
    687         data = {
    688             model: self._make_maximum_spares(counts)
    689                 for model, counts in self._model_data
    690         }
    691         inventory = create_inventory(data)
    692         message = lab_inventory._generate_model_inventory_message(inventory)
    693         self._check_message(message)
    694 
    695 
    696     def test_ignore_no_spares(self):
    697         """Test that messages ignore models with no spare pool."""
    698         data = {
    699             model: self._make_maximum_spares(counts)
    700                 for model, counts in self._model_data
    701         }
    702         data['elephant'] = ((5, 4, 0), (0, 0, 0))
    703         inventory = create_inventory(data)
    704         message = lab_inventory._generate_model_inventory_message(inventory)
    705         self._check_message(message)
    706 
    707 
    708     def test_ignore_no_critical(self):
    709         """Test that messages ignore models with no critical pools."""
    710         data = {
    711             model: self._make_maximum_spares(counts)
    712                 for model, counts in self._model_data
    713         }
    714         data['elephant'] = ((0, 0, 0), (1, 5, 1))
    715         inventory = create_inventory(data)
    716         message = lab_inventory._generate_model_inventory_message(inventory)
    717         self._check_message(message)
    718 
    719 
    720     def test_ignore_no_bad(self):
    721         """Test that messages ignore models with no bad DUTs."""
    722         data = {
    723             model: self._make_maximum_spares(counts)
    724                 for model, counts in self._model_data
    725         }
    726         data['elephant'] = ((5, 0, 1), (5, 0, 1))
    727         inventory = create_inventory(data)
    728         message = lab_inventory._generate_model_inventory_message(inventory)
    729         self._check_message(message)
    730 
    731 
    732 class _PoolInventoryTestBase(unittest.TestCase):
    733     """Parent class for tests relating to generating pool inventory messages.
    734 
    735     Func `setUp` in the class parses a given |message_template| to obtain
    736     header and body.
    737     """
    738     def _read_template(self, message_template):
    739         """Read message template for PoolInventoryTest and IdleInventoryTest.
    740 
    741         @param message_template: the input template to be parsed into: header
    742         and content (report_lines).
    743 
    744         """
    745         message_lines = message_template.split('\n')
    746         self._header = message_lines[1]
    747         self._report_lines = message_lines[2:-1]
    748 
    749 
    750     def _check_report_no_info(self, text):
    751         """Test a message body containing no reported info.
    752 
    753         The input `text` was created from a query to an inventory, which has
    754         no objects meet the query and leads to an `empty` return. Assert that
    755         the text consists of a single line starting with '(' and ending with ')'.
    756 
    757         @param text: Message body text to be tested.
    758 
    759         """
    760         self.assertTrue(len(text) == 1 and
    761                             text[0][0] == '(' and
    762                             text[0][-1] == ')')
    763 
    764 
    765     def _check_report(self, text):
    766         """Test a message against the passed |expected_content|.
    767 
    768         @param text: Message body text to be tested.
    769         @param expected_content: The ground-truth content to be compared with.
    770 
    771         """
    772         self.assertEqual(text, self._report_lines)
    773 
    774 
    775 # _POOL_MESSAGE_TEMPLATE -
    776 # This is a sample of the output text produced by
    777 # _generate_pool_inventory_message().  This string is parsed by the
    778 # tests below to construct a sample inventory that should produce
    779 # the output, and then the output is generated and checked against
    780 # this original sample.
    781 #
    782 # See the comments on _BOARD_MESSAGE_TEMPLATE above for the
    783 # rationale on using sample text in this way.
    784 
    785 _POOL_MESSAGE_TEMPLATE = '''
    786 Model                    Bad  Idle  Good Total
    787 lion                       5     2     6    13
    788 tiger                      4     1     5    10
    789 bear                       3     0     7    10
    790 aardvark                   2     0     0     2
    791 platypus                   1     1     1     3
    792 '''
    793 
    794 _POOL_ADMIN_URL = 'http://go/cros-manage-duts'
    795 
    796 
    797 class PoolInventoryTests(_PoolInventoryTestBase):
    798     """Tests for `_generate_pool_inventory_message()`.
    799 
    800     The tests create various test inventories designed to match the
    801     counts in `_POOL_MESSAGE_TEMPLATE`, and assert that the
    802     generated message text matches the format established in the
    803     original message text.
    804 
    805     The output message text is parsed against the following grammar:
    806         <message> -> <intro> <pool> { "blank line" <pool> }
    807         <intro> ->
    808             Instructions to depty mentioning the admin page URL
    809             A blank line
    810         <pool> ->
    811             <description>
    812             <header line>
    813             <message body>
    814         <description> ->
    815             Any number of lines describing one pool
    816         <header line> ->
    817             The header line from `_POOL_MESSAGE_TEMPLATE`
    818         <message body> ->
    819             Any number of non-blank lines
    820 
    821     After parsing messages into the parts described above, various
    822     assertions are tested against the parsed output, including
    823     that the message body matches the body from
    824     `_POOL_MESSAGE_TEMPLATE`.
    825 
    826     Parse message text is represented as a list of strings, split on
    827     the `'\n'` separator.
    828     """
    829 
    830     def setUp(self):
    831         super(PoolInventoryTests, self)._read_template(_POOL_MESSAGE_TEMPLATE)
    832         self._model_data = []
    833         for l in self._report_lines:
    834             items = l.split()
    835             model = items[0]
    836             bad = int(items[1])
    837             idle = int(items[2])
    838             good = int(items[3])
    839             self._model_data.append((model, (good, bad, idle)))
    840 
    841 
    842     def _create_histories(self, pools, model_data):
    843         """Return a list suitable to create a `_LabInventory` object.
    844 
    845         Creates a list of `_FakeHostHistory` objects that can be
    846         used to create a lab inventory.  `pools` is a list of strings
    847         naming pools, and `model_data` is a list of tuples of the
    848         form
    849             `(model, (goodcount, badcount))`
    850         where
    851             `model` is a model name.
    852             `goodcount` is the number of working DUTs in the pool.
    853             `badcount` is the number of broken DUTs in the pool.
    854 
    855         @param pools       List of pools for which to create
    856                            histories.
    857         @param model_data  List of tuples containing models and DUT
    858                            counts.
    859         @return A list of `_FakeHostHistory` objects that can be
    860                 used to create a `_LabInventory` object.
    861 
    862         """
    863         histories = []
    864         status_choices = (_WORKING, _BROKEN, _UNUSED)
    865         for pool in pools:
    866             for model, counts in model_data:
    867                 for status, count in zip(status_choices, counts):
    868                     for x in range(0, count):
    869                         histories.append(
    870                             _FakeHostHistory(model, pool, status))
    871         return histories
    872 
    873 
    874     def _parse_pool_summaries(self, histories):
    875         """Parse message output according to the grammar above.
    876 
    877         Create a lab inventory from the given `histories`, and
    878         generate the pool inventory message.  Then parse the message
    879         and return a dictionary mapping each pool to the message
    880         body parsed after that pool.
    881 
    882         Tests the following assertions:
    883           * Each <description> contains a mention of exactly one
    884             pool in the `CRITICAL_POOLS` list.
    885           * Each pool is mentioned in exactly one <description>.
    886         Note that the grammar requires the header to appear once
    887         for each pool, so the parsing implicitly asserts that the
    888         output contains the header.
    889 
    890         @param histories  Input used to create the test
    891                           `_LabInventory` object.
    892         @return A dictionary mapping model names to the output
    893                 (a list of lines) for the model.
    894 
    895         """
    896         inventory = lab_inventory._LabInventory(
    897                 histories, lab_inventory.MANAGED_POOLS)
    898         message = lab_inventory._generate_pool_inventory_message(
    899                 inventory).split('\n')
    900         poolset = set(lab_inventory.CRITICAL_POOLS)
    901         seen_url = False
    902         seen_intro = False
    903         description = ''
    904         model_text = {}
    905         current_pool = None
    906         for line in message:
    907             if not seen_url:
    908                 if _POOL_ADMIN_URL in line:
    909                     seen_url = True
    910             elif not seen_intro:
    911                 if not line:
    912                     seen_intro = True
    913             elif current_pool is None:
    914                 if line == self._header:
    915                     pools_mentioned = [p for p in poolset
    916                                            if p in description]
    917                     self.assertEqual(len(pools_mentioned), 1)
    918                     current_pool = pools_mentioned[0]
    919                     description = ''
    920                     model_text[current_pool] = []
    921                     poolset.remove(current_pool)
    922                 else:
    923                     description += line
    924             else:
    925                 if line:
    926                     model_text[current_pool].append(line)
    927                 else:
    928                     current_pool = None
    929         self.assertEqual(len(poolset), 0)
    930         return model_text
    931 
    932 
    933     def test_no_shortages(self):
    934         """Test correct output when no pools have shortages."""
    935         model_text = self._parse_pool_summaries([])
    936         for text in model_text.values():
    937             self._check_report_no_info(text)
    938 
    939 
    940     def test_one_pool_shortage(self):
    941         """Test correct output when exactly one pool has a shortage."""
    942         for pool in lab_inventory.CRITICAL_POOLS:
    943             histories = self._create_histories((pool,),
    944                                                self._model_data)
    945             model_text = self._parse_pool_summaries(histories)
    946             for checkpool in lab_inventory.CRITICAL_POOLS:
    947                 text = model_text[checkpool]
    948                 if checkpool == pool:
    949                     self._check_report(text)
    950                 else:
    951                     self._check_report_no_info(text)
    952 
    953 
    954     def test_all_pool_shortages(self):
    955         """Test correct output when all pools have a shortage."""
    956         histories = []
    957         for pool in lab_inventory.CRITICAL_POOLS:
    958             histories.extend(
    959                 self._create_histories((pool,),
    960                                        self._model_data))
    961         model_text = self._parse_pool_summaries(histories)
    962         for pool in lab_inventory.CRITICAL_POOLS:
    963             self._check_report(model_text[pool])
    964 
    965 
    966     def test_full_model_ignored(self):
    967         """Test that models at full strength are not reported."""
    968         pool = lab_inventory.CRITICAL_POOLS[0]
    969         full_model = [('echidna', (5, 0, 0))]
    970         histories = self._create_histories((pool,),
    971                                            full_model)
    972         text = self._parse_pool_summaries(histories)[pool]
    973         self._check_report_no_info(text)
    974         model_data = self._model_data + full_model
    975         histories = self._create_histories((pool,), model_data)
    976         text = self._parse_pool_summaries(histories)[pool]
    977         self._check_report(text)
    978 
    979 
    980     def test_spare_pool_ignored(self):
    981         """Test that reporting ignores the spare pool inventory."""
    982         spare_pool = lab_inventory.SPARE_POOL
    983         spare_data = self._model_data + [('echidna', (0, 5, 0))]
    984         histories = self._create_histories((spare_pool,),
    985                                            spare_data)
    986         model_text = self._parse_pool_summaries(histories)
    987         for pool in lab_inventory.CRITICAL_POOLS:
    988             self._check_report_no_info(model_text[pool])
    989 
    990 
    991 _IDLE_MESSAGE_TEMPLATE = '''
    992 Hostname                       Model                Pool
    993 chromeos4-row12-rack4-host7    tiger                bvt
    994 chromeos1-row3-rack1-host2     lion                 bvt
    995 chromeos3-row2-rack2-host5     lion                 cq
    996 chromeos2-row7-rack3-host11    platypus             suites
    997 '''
    998 
    999 
   1000 class IdleInventoryTests(_PoolInventoryTestBase):
   1001     """Tests for `_generate_idle_inventory_message()`.
   1002 
   1003     The tests create idle duts that match the counts and pool in
   1004     `_IDLE_MESSAGE_TEMPLATE`. In test, it asserts that the generated
   1005     idle message text matches the format established in
   1006     `_IDLE_MESSAGE_TEMPLATE`.
   1007 
   1008     Parse message text is represented as a list of strings, split on
   1009     the `'\n'` separator.
   1010 
   1011     """
   1012 
   1013     def setUp(self):
   1014         super(IdleInventoryTests, self)._read_template(_IDLE_MESSAGE_TEMPLATE)
   1015         self._host_data = []
   1016         for h in self._report_lines:
   1017             items = h.split()
   1018             hostname = items[0]
   1019             model = items[1]
   1020             pool = items[2]
   1021             self._host_data.append((hostname, model, pool))
   1022         self._histories = []
   1023         self._histories.append(_FakeHostHistory('echidna', 'bvt', _BROKEN))
   1024         self._histories.append(_FakeHostHistory('lion', 'bvt', _WORKING))
   1025 
   1026 
   1027     def _add_idles(self):
   1028         """Add idle duts from `_IDLE_MESSAGE_TEMPLATE`."""
   1029         idle_histories = [_FakeHostHistory(
   1030                 model, pool, _UNUSED, hostname)
   1031                         for hostname, model, pool in self._host_data]
   1032         self._histories.extend(idle_histories)
   1033 
   1034 
   1035     def _check_header(self, text):
   1036         """Check whether header in the template `_IDLE_MESSAGE_TEMPLATE` is in
   1037         passed text."""
   1038         self.assertIn(self._header, text)
   1039 
   1040 
   1041     def _get_idle_message(self, histories):
   1042         """Generate idle inventory and obtain its message.
   1043 
   1044         @param histories: Used to create lab inventory.
   1045 
   1046         @return the generated idle message.
   1047 
   1048         """
   1049         inventory = lab_inventory._LabInventory(
   1050                 histories, lab_inventory.MANAGED_POOLS)
   1051         message = lab_inventory._generate_idle_inventory_message(
   1052                 inventory).split('\n')
   1053         return message
   1054 
   1055 
   1056     def test_check_idle_inventory(self):
   1057         """Test that reporting all the idle DUTs for every pool, sorted by
   1058         lab_inventory.MANAGED_POOLS.
   1059         """
   1060         self._add_idles()
   1061 
   1062         message = self._get_idle_message(self._histories)
   1063         self._check_header(message)
   1064         self._check_report(message[message.index(self._header) + 1 :])
   1065 
   1066 
   1067     def test_no_idle_inventory(self):
   1068         """Test that reporting no idle DUTs."""
   1069         message = self._get_idle_message(self._histories)
   1070         self._check_header(message)
   1071         self._check_report_no_info(
   1072                 message[message.index(self._header) + 1 :])
   1073 
   1074 
   1075 class CommandParsingTests(unittest.TestCase):
   1076     """Tests for command line argument parsing in `_parse_command()`."""
   1077 
   1078     # At least one of these options must be specified on every command
   1079     # line; otherwise, the command line parsing will fail.
   1080     _REPORT_OPTIONS = ['--model-notify=', '--pool-notify=', '--repair-loops']
   1081 
   1082     def setUp(self):
   1083         dirpath = '/usr/local/fubar'
   1084         self._command_path = os.path.join(dirpath,
   1085                                           'site_utils',
   1086                                           'arglebargle')
   1087         self._logdir = os.path.join(dirpath, lab_inventory._LOGDIR)
   1088 
   1089 
   1090     def _parse_arguments(self, argv):
   1091         """Test parsing with explictly passed report options."""
   1092         full_argv = [self._command_path] + argv
   1093         return lab_inventory._parse_command(full_argv)
   1094 
   1095 
   1096     def _parse_non_report_arguments(self, argv):
   1097         """Test parsing for non-report command-line options."""
   1098         return self._parse_arguments(argv + self._REPORT_OPTIONS)
   1099 
   1100 
   1101     def _check_non_report_defaults(self, report_option):
   1102         arguments = self._parse_arguments([report_option])
   1103         self.assertEqual(arguments.duration,
   1104                          lab_inventory._DEFAULT_DURATION)
   1105         self.assertIsNone(arguments.recommend)
   1106         self.assertFalse(arguments.debug)
   1107         self.assertEqual(arguments.logdir, self._logdir)
   1108         self.assertEqual(arguments.modelnames, [])
   1109         return arguments
   1110 
   1111 
   1112     def test_empty_arguments(self):
   1113         """Test that no reports requested is an error."""
   1114         arguments = self._parse_arguments([])
   1115         self.assertIsNone(arguments)
   1116 
   1117 
   1118     def test_argument_defaults(self):
   1119         """Test that option defaults match expectations."""
   1120         for report in self._REPORT_OPTIONS:
   1121             arguments = self._check_non_report_defaults(report)
   1122 
   1123 
   1124     def test_model_notify_defaults(self):
   1125         """Test defaults when `--model-notify` is specified alone."""
   1126         arguments = self._parse_arguments(['--model-notify='])
   1127         self.assertEqual(arguments.model_notify, [''])
   1128         self.assertEqual(arguments.pool_notify, [])
   1129         self.assertFalse(arguments.repair_loops)
   1130 
   1131 
   1132     def test_pool_notify_defaults(self):
   1133         """Test defaults when `--pool-notify` is specified alone."""
   1134         arguments = self._parse_arguments(['--pool-notify='])
   1135         self.assertEqual(arguments.model_notify, [])
   1136         self.assertEqual(arguments.pool_notify, [''])
   1137         self.assertFalse(arguments.repair_loops)
   1138 
   1139 
   1140     def test_repair_loop_defaults(self):
   1141         """Test defaults when `--repair-loops` is specified alone."""
   1142         arguments = self._parse_arguments(['--repair-loops'])
   1143         self.assertEqual(arguments.model_notify, [])
   1144         self.assertEqual(arguments.pool_notify, [])
   1145         self.assertTrue(arguments.repair_loops)
   1146 
   1147 
   1148     def test_model_arguments(self):
   1149         """Test that non-option arguments are returned in `modelnames`."""
   1150         modellist = ['aardvark', 'echidna']
   1151         arguments = self._parse_non_report_arguments(modellist)
   1152         self.assertEqual(arguments.modelnames, modellist)
   1153 
   1154 
   1155     def test_recommend_option(self):
   1156         """Test parsing of the `--recommend` option."""
   1157         for opt in ['-r', '--recommend']:
   1158             for recommend in ['5', '55']:
   1159                 arguments = self._parse_non_report_arguments([opt, recommend])
   1160                 self.assertEqual(arguments.recommend, int(recommend))
   1161 
   1162 
   1163     def test_debug_option(self):
   1164         """Test parsing of the `--debug` option."""
   1165         arguments = self._parse_non_report_arguments(['--debug'])
   1166         self.assertTrue(arguments.debug)
   1167 
   1168 
   1169     def test_duration(self):
   1170         """Test parsing of the `--duration` option."""
   1171         for opt in ['-d', '--duration']:
   1172             for duration in ['1', '11']:
   1173                 arguments = self._parse_non_report_arguments([opt, duration])
   1174                 self.assertEqual(arguments.duration, int(duration))
   1175 
   1176 
   1177     def _check_email_option(self, option, getlist):
   1178         """Test parsing of e-mail address options.
   1179 
   1180         This is a helper function to test the `--model-notify` and
   1181         `--pool-notify` options.  It tests the following cases:
   1182           * `--option a1` gives the list [a1]
   1183           * `--option ' a1 '` gives the list [a1]
   1184           * `--option a1 --option a2` gives the list [a1, a2]
   1185           * `--option a1,a2` gives the list [a1, a2]
   1186           * `--option 'a1, a2'` gives the list [a1, a2]
   1187 
   1188         @param option  The option to be tested.
   1189         @param getlist A function to return the option's value from
   1190                        parsed command line arguments.
   1191 
   1192         """
   1193         a1 = 'mumble (at] mumbler.com'
   1194         a2 = 'bumble (at] bumbler.org'
   1195         arguments = self._parse_arguments([option, a1])
   1196         self.assertEqual(getlist(arguments), [a1])
   1197         arguments = self._parse_arguments([option, ' ' + a1 + ' '])
   1198         self.assertEqual(getlist(arguments), [a1])
   1199         arguments = self._parse_arguments([option, a1, option, a2])
   1200         self.assertEqual(getlist(arguments), [a1, a2])
   1201         arguments = self._parse_arguments(
   1202                 [option, ','.join([a1, a2])])
   1203         self.assertEqual(getlist(arguments), [a1, a2])
   1204         arguments = self._parse_arguments(
   1205                 [option, ', '.join([a1, a2])])
   1206         self.assertEqual(getlist(arguments), [a1, a2])
   1207 
   1208 
   1209     def test_model_notify(self):
   1210         """Test parsing of the `--model-notify` option."""
   1211         self._check_email_option('--model-notify',
   1212                                  lambda a: a.model_notify)
   1213 
   1214 
   1215     def test_pool_notify(self):
   1216         """Test parsing of the `--pool-notify` option."""
   1217         self._check_email_option('--pool-notify',
   1218                                  lambda a: a.pool_notify)
   1219 
   1220 
   1221     def test_logdir_option(self):
   1222         """Test parsing of the `--logdir` option."""
   1223         logdir = '/usr/local/whatsis/logs'
   1224         arguments = self._parse_non_report_arguments(['--logdir', logdir])
   1225         self.assertEqual(arguments.logdir, logdir)
   1226 
   1227 
   1228 if __name__ == '__main__':
   1229     # Some of the functions we test log messages.  Prevent those
   1230     # messages from showing up in test output.
   1231     logging.getLogger().setLevel(logging.CRITICAL)
   1232     unittest.main()
   1233