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