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