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 collections 7import itertools 8import logging 9import os 10import unittest 11 12import common 13from autotest_lib.server.lib import status_history 14from autotest_lib.site_utils import lab_inventory 15 16 17class _FakeHost(object): 18 """Class to mock `Host` in _FakeHostHistory for testing.""" 19 20 def __init__(self, hostname): 21 self.hostname = hostname 22 23 24class _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 40class _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 67class 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 145class 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 265class 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 325class 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. 429STATUS_CHOICES = (_WORKING, _BROKEN, _UNUSED) 430StatusCounts = 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. 433CRITICAL_POOL = lab_inventory.CRITICAL_POOLS[0] 434SPARE_POOL = lab_inventory.SPARE_POOL 435POOL_CHOICES = (CRITICAL_POOL, SPARE_POOL) 436PoolStatusCounts = collections.namedtuple('PoolStatusCounts', 437 ['critical', 'spare']) 438 439def 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 463class 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 = ''' 593Model Avail Bad Idle Good Spare Total 594lion -1 13 2 11 12 26 595tiger -1 5 2 9 4 16 596bear 0 5 2 10 5 17 597platypus 4 2 2 20 6 24 598aardvark 7 2 2 6 9 10 599''' 600 601 602class 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 732class _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 = ''' 786Model Bad Idle Good Total 787lion 5 2 6 13 788tiger 4 1 5 10 789bear 3 0 7 10 790aardvark 2 0 0 2 791platypus 1 1 1 3 792''' 793 794_POOL_ADMIN_URL = 'http://go/cros-manage-duts' 795 796 797class 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 = ''' 992Hostname Model Pool 993chromeos4-row12-rack4-host7 tiger bvt 994chromeos1-row3-rack1-host2 lion bvt 995chromeos3-row2-rack2-host5 lion cq 996chromeos2-row7-rack3-host11 platypus suites 997''' 998 999 1000class 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 1075class 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@mumbler.com' 1194 a2 = 'bumble@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 1228if __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