Support lab inventory summary by model

LabInventory can now summarize DUT histories by board as well as model.
This will allow all scripts that use LabInventory to slowly migrate to
respecting models instead of boards.

BUG=chromium:780892
TEST=unittests.

Change-Id: I422b4bd151b3bad97591eed3da23b6c12c968414
Reviewed-on: https://chromium-review.googlesource.com/760498
Reviewed-by: Prathmesh Prabhu <pprabhu@chromium.org>
Commit-Queue: Prathmesh Prabhu <pprabhu@chromium.org>
Tested-by: Prathmesh Prabhu <pprabhu@chromium.org>
diff --git a/site_utils/lab_inventory.py b/site_utils/lab_inventory.py
index f538255..2fc0b44 100755
--- a/site_utils/lab_inventory.py
+++ b/site_utils/lab_inventory.py
@@ -481,7 +481,9 @@
         self.histories = histories
         self._dut_count = len(histories)
         self._managed_boards = {}
+        self._managed_models = {}
         self.by_board = self._classify_by_label_type('board')
+        self.by_model = self._classify_by_label_type('model')
 
 
     def _classify_by_label_type(self, label_key):
@@ -501,6 +503,38 @@
     def get_managed_boards(self, pool=_MANAGED_POOL_DEFAULT):
         """Return the set of "managed" boards.
 
+        @param pool: The specified pool for managed boards.
+        @return A set of all the boards that have both spare and
+                non-spare pools, unless the pool is specified,
+                then the set of boards in that pool.
+        """
+        if self._managed_boards.get(pool) is None:
+            self._managed_boards[pool] = set()
+            for board, counts in self.by_board.iteritems():
+                if self._is_managed(pool, counts):
+                    self._managed_boards[pool].add(board)
+        return self._managed_boards[pool]
+
+
+    def get_managed_models(self, pool=_MANAGED_POOL_DEFAULT):
+        """Return the set of "managed" models.
+
+        @param pool: The specified pool for managed models.
+        @return A set of all the models that have both spare and
+                non-spare pools, unless the pool is specified,
+                then the set of models in that pool.
+        """
+        if self._managed_models.get(pool) is None:
+            self._managed_models[pool] = set()
+            for board, counts in self.by_model.iteritems():
+                if self._is_managed(pool, counts):
+                    self._managed_models[pool].add(board)
+        return self._managed_models[pool]
+
+
+    def _is_managed(self, pool, histories):
+        """Deterime if the given histories contain DUTs to be managed for pool.
+
         Operationally, saying a board is "managed" means that the
         board will be included in the "board" and "repair
         recommendations" reports.  That is, if there are failures in
@@ -511,25 +545,15 @@
         has DUTs in both the spare and a non-spare (i.e. critical)
         pool.
 
-        @param pool: The specified pool for managed boards.
-        @return A set of all the boards that have both spare and
-                non-spare pools, unless the pool is specified,
-                then the set of boards in that pool.
         """
-        if self._managed_boards.get(pool, None) is None:
-            self._managed_boards[pool] = set()
-            for board, counts in self.by_board.iteritems():
-                # Get the counts for all pools, otherwise get it for the
-                # specified pool.
-                if pool == _MANAGED_POOL_DEFAULT:
-                    spares = counts.get_total(SPARE_POOL)
-                    total = counts.get_total()
-                    if spares != 0 and spares != total:
-                        self._managed_boards[pool].add(board)
-                else:
-                    if counts.get_total(pool) != 0:
-                        self._managed_boards[pool].add(board)
-        return self._managed_boards[pool]
+        # Get the counts for all pools, otherwise get it for the
+        # specified pool.
+        if pool == _MANAGED_POOL_DEFAULT:
+            spares = histories.get_total(SPARE_POOL)
+            total = histories.get_total()
+            return spares != 0 and spares != total
+        else:
+            return histories.get_total(pool) != 0
 
 
     def get_num_duts(self):
@@ -542,6 +566,11 @@
         return len(self.by_board)
 
 
+    def get_num_models(self):
+        """Return the total number of models in the inventory."""
+        return len(self.by_model)
+
+
 def _sort_by_location(inventory_list):
     """Return a list of DUTs, organized by location.
 
diff --git a/site_utils/lab_inventory_unittest.py b/site_utils/lab_inventory_unittest.py
index ae5b317..9b956d5 100755
--- a/site_utils/lab_inventory_unittest.py
+++ b/site_utils/lab_inventory_unittest.py
@@ -505,9 +505,7 @@
 
     """
 
-    # _BOARD_LIST - A list of sample board names for use in testing.
-
-    _BOARD_LIST = [
+    _BOARD_OR_MODEL_LIST = [
         'lion',
         'tiger',
         'bear',
@@ -519,7 +517,8 @@
     ]
 
 
-    def _check_inventory_details(self, inventory, data, by_board=True):
+    def _check_inventory_details(self, inventory, data, by_board=True,
+                                 msg=None):
         """Some common detailed inventory checks.
 
         The checks here are common to many tests below. At the same time, thsese
@@ -552,14 +551,17 @@
                             histories.get_idle(SPARE_POOL),
                     ),
             )
-            self.assertEqual(data[key], calculated_counts)
+            self.assertEqual(data[key], calculated_counts, msg)
 
             self.assertEqual(len(histories.get_working_list()),
-                             sum([p.good for p in data[key]]))
+                             sum([p.good for p in data[key]]),
+                             msg)
             self.assertEqual(len(histories.get_broken_list()),
-                             sum([p.bad for p in data[key]]))
+                             sum([p.bad for p in data[key]]),
+                             msg)
             self.assertEqual(len(histories.get_idle_list()),
-                             sum([p.unused for p in data[key]]))
+                             sum([p.unused for p in data[key]]),
+                             msg)
 
 
     def test_empty(self):
@@ -568,7 +570,10 @@
         self.assertEqual(inventory.get_num_duts(), 0)
         self.assertEqual(inventory.get_num_boards(), 0)
         self.assertEqual(inventory.get_managed_boards(), set())
-        self._check_inventory_details(inventory, {})
+        self._check_inventory_details(inventory, {}, by_board=True)
+        self.assertEqual(inventory.get_num_models(), 0)
+        self.assertEqual(inventory.get_managed_models(), set())
+        self._check_inventory_details(inventory, {}, by_board=False)
 
 
     def test_missing_board(self):
@@ -582,20 +587,23 @@
         self.assertEqual(inventory.get_num_duts(), 0)
         self.assertEqual(inventory.get_num_boards(), 0)
         self.assertEqual(inventory.get_managed_boards(), set())
-        self._check_inventory_details(inventory, {})
+        self._check_inventory_details(inventory, {}, by_board=True)
+        self.assertEqual(inventory.get_num_models(), 0)
+        self.assertEqual(inventory.get_managed_models(), set())
+        self._check_inventory_details(inventory, {}, by_board=False)
 
 
     def test_board_counts(self):
         """Test counts for various numbers of boards."""
-        for board_count in [1, 2, len(self._BOARD_LIST)]:
+        for board_count in [1, 2, len(self._BOARD_OR_MODEL_LIST)]:
             self.parameterized_test_board_count(board_count)
 
 
     def parameterized_test_board_count(self, board_count):
         """Parameterized test for testing a specific number of boards."""
         self.longMessage = True
-        msg = '[board_count: %s]' % (board_count,)
-        boards = self._BOARD_LIST[:board_count]
+        msg = '[board_count: %s]' % (board_count)
+        boards = self._BOARD_OR_MODEL_LIST[:board_count]
         data = {
                 b: PoolStatusCounts(
                         StatusCounts(1, 1, 1),
@@ -603,13 +611,41 @@
                 )
                 for b in boards
         }
-        inventory = create_inventory(data)
+        inventory = create_inventory(data, by_board=True)
         self.assertEqual(inventory.get_num_duts(), 6 * board_count, msg)
         self.assertEqual(inventory.get_num_boards(), board_count, msg)
         self.assertEqual(inventory.get_managed_boards(), set(boards), msg)
-        self._check_inventory_details(inventory, data, msg)
+        self._check_inventory_details(inventory, data, by_board=True, msg=msg)
+        self.assertEqual(inventory.get_num_models(), 1, msg)
+        self.assertEqual(inventory.get_managed_models(), {'dummy_model'}, msg)
 
 
+    def test_model_counts(self):
+        """Test counts for various numbers of models."""
+        for model_count in [1, 2, len(self._BOARD_OR_MODEL_LIST)]:
+            self.parameterized_test_model_count(model_count)
+
+
+    def parameterized_test_model_count(self, model_count):
+        """Parameterized test for testing a specific number of models."""
+        self.longMessage = True
+        msg = '[model: %s]' % (model_count)
+        models = self._BOARD_OR_MODEL_LIST[:model_count]
+        data = {
+                m: PoolStatusCounts(
+                        StatusCounts(1, 1, 1),
+                        StatusCounts(1, 1, 1),
+                )
+                for m in models
+        }
+        inventory = create_inventory(data, by_board=False)
+        self.assertEqual(inventory.get_num_duts(), 6 * model_count, msg)
+        self.assertEqual(inventory.get_num_models(), model_count, msg)
+        self.assertEqual(inventory.get_managed_models(), set(models), msg)
+        self._check_inventory_details(inventory, data, by_board=False, msg=msg)
+        self.assertEqual(inventory.get_num_boards(), 1, msg)
+        self.assertEqual(inventory.get_managed_boards(), {'dummy_board'}, msg)
+
 
     def test_single_dut_counts(self):
         """Test counts when there is a single DUT per board, and it is good."""
@@ -626,14 +662,14 @@
     def parameterized_test_single_dut_counts(self, counts):
         """Parmeterized test for single dut counts."""
         self.longMessage = True
-        board = self._BOARD_LIST[0]
+        board = self._BOARD_OR_MODEL_LIST[0]
         data = {board: counts}
         msg = '[data: %s]' % (data,)
         inventory = create_inventory(data)
         self.assertEqual(inventory.get_num_duts(), 1, msg)
         self.assertEqual(inventory.get_num_boards(), 1, msg)
         self.assertEqual(inventory.get_managed_boards(), set(), msg)
-        self._check_inventory_details(inventory, data, msg)
+        self._check_inventory_details(inventory, data, by_board=True, msg=msg)
 
 
 # BOARD_MESSAGE_TEMPLATE -