mobmonitor: Small additions to the disk info object.

This CL updates some object names and return types for
collecting information about block devices. Block device
size is now added to the collection.

BUG=chromium:522731
TEST=Unittests and manually on moblab.

Change-Id: Ic5dd8bbc45aa637bf38afc4bcc64b025e1cf04e8
Reviewed-on: https://chromium-review.googlesource.com/296454
Commit-Ready: Matthew Sartori <msartori@chromium.org>
Tested-by: Matthew Sartori <msartori@chromium.org>
Reviewed-by: Paul Hobbs <phobbs@google.com>
Reviewed-by: Simran Basi <sbasi@chromium.org>

3 files changed, 66 insertions, 33 deletions

diff --git a/mobmonitor/checkfile/manager.py b/mobmonitor/checkfile/manager.py
index fe20d25d2..4af79e4e8 100644
--- a/mobmonitor/checkfile/manager.py
+++ b/mobmonitor/checkfile/manager.py
@@ -483,7 +483,7 @@ class CheckFileManager(object):
         # all services.
         self.Execute(force=True)
         self.ConsolidateServiceStates()
-      except TypeError, e:
+      except Exception, e:
         LOGGER.error('Failed to execute the repair action "%s"'
                      ' for service "%s": %s', action, service, e,
                      exc_info=True)
diff --git a/mobmonitor/system/systeminfo.py b/mobmonitor/system/systeminfo.py
index 90dcfa75a..71b8ad4ea 100644
--- a/mobmonitor/system/systeminfo.py
+++ b/mobmonitor/system/systeminfo.py
@@ -42,6 +42,7 @@ import os
 import time
 
 from chromite.lib import cros_build_lib
+from chromite.lib import osutils
 
 
 SYSTEMFILE_PROC_MOUNTS = '/proc/mounts'
@@ -62,7 +63,7 @@ UPDATE_CPU_SEC = 2
 RESOURCENAME_MEMORY = 'memory'
 RESOURCENAME_DISKPARTITIONS = 'diskpartitions'
 RESOURCENAME_DISKUSAGE = 'diskusage'
-RESOURCENAME_DISKBY = 'diskby'
+RESOURCENAME_BLOCKDEVICE = 'blockdevice'
 RESOURCENAME_CPUPREVTIMES = 'cpuprevtimes'
 RESOURCENAME_CPUTIMES = 'cputimes'
 RESOURCENAME_CPULOADS = 'cpuloads'
@@ -75,7 +76,9 @@ RESOURCE_DISKPARTITION = collections.namedtuple('diskpartition',
 RESOURCE_DISKUSAGE = collections.namedtuple('diskusage', ['total', 'used',
                                                           'free',
                                                           'percent_used'])
-RESOURCE_DISKBY = collections.namedtuple('diskby', ['device', 'ids', 'labels'])
+RESOURCE_BLOCKDEVICE = collections.namedtuple('blockdevice',
+                                              ['device', 'size', 'ids',
+                                               'labels'])
 RESOURCE_CPUTIME = collections.namedtuple('cputime', ['cpu', 'total', 'idle',
                                                       'nonidle'])
 RESOURCE_CPULOAD = collections.namedtuple('cpuload', ['cpu', 'load'])
@@ -331,40 +334,62 @@ class Disk(SystemInfoStorage):
 
     return diskusage
 
-  @CheckStorage(RESOURCENAME_DISKBY)
-  def DiskBy(self, device):
-    """Collects information under the SYSTEMFILE_DEV_DISKBY directories.
+  @CheckStorage(RESOURCENAME_BLOCKDEVICE)
+  def BlockDevices(self, device=''):
+    """Collects information about block devices.
 
-    Args:
-      device: This is the same as the 'devicename' attribute given
-        in the return value of DiskPartitions.
+    This method combines information from:
+      (1) Reading through the SYSTEMFILE_DEV_DISKBY directories.
+      (2) Executing the 'lsblk' command provided by osutils.ListBlockDevices.
 
     Returns:
-      A named tuple with the following fields:
-        device: The name of the device, same as the given argument.
+      A list of named tuples. Each tuple has the following fields:
+        device: The name of the block device.
+        size: The size of the block device in bytes.
         ids: A list of ids assigned to this device.
         labels: A list of labels assigned to this device.
     """
-    ids = []
-    labels = []
+    devicefilter = os.path.basename(device)
+
+    # Data collected from the SYSTEMFILE_DEV_DISKBY directories.
+    ids = {}
+    labels = {}
 
-    devicename = os.path.basename(device)
+    # Data collected from 'lsblk'.
+    sizes = {}
+
+    # Collect diskby information.
     for prop, diskdir in SYSTEMFILE_DEV_DISKBY.iteritems():
-      cmd = ['find', diskdir, '-lname', '*%s' % devicename]
+      cmd = ['find', diskdir, '-lname', '*%s' % devicefilter]
       cmd_result = cros_build_lib.RunCommand(cmd, log_output=True)
 
       if not cmd_result.output:
         continue
 
       results = cmd_result.output.split()
-      results = [os.path.basename(p) for p in results]
-
-      if 'ids' == prop:
-        ids = results
-      elif 'labels' == prop:
-        labels = results
-
-    return RESOURCE_DISKBY(device, ids, labels)
+      for result in results:
+        devicename = os.path.abspath(osutils.ResolveSymlink(result))
+        result = os.path.basename(result)
+
+        # Ensure that each of our data dicts have the same keys.
+        ids.setdefault(devicename, [])
+        labels.setdefault(devicename, [])
+        sizes.setdefault(devicename, 0)
+
+        if 'ids' == prop:
+          ids[devicename].append(result)
+        elif 'labels' == prop:
+          labels[devicename].append(result)
+
+    # Collect lsblk information.
+    for device in osutils.ListBlockDevices(in_bytes=True):
+      devicename = os.path.join('/dev', device.NAME)
+      if devicename in ids:
+        sizes[devicename] = int(device.SIZE)
+
+    return [RESOURCE_BLOCKDEVICE(device, sizes[device], ids[device],
+                                 labels[device])
+            for device in ids.iterkeys()]
 
 
 class Cpu(SystemInfoStorage):
diff --git a/mobmonitor/system/systeminfo_unittest.py b/mobmonitor/system/systeminfo_unittest.py
index 04f3deb24..f92b4fd41 100644
--- a/mobmonitor/system/systeminfo_unittest.py
+++ b/mobmonitor/system/systeminfo_unittest.py
@@ -14,6 +14,7 @@ import time
 
 from chromite.lib import cros_build_lib_unittest
 from chromite.lib import cros_test_lib
+from chromite.lib import osutils
 from chromite.lib import partial_mock
 from chromite.mobmonitor.system import systeminfo
 
@@ -344,27 +345,31 @@ class DiskTest(cros_test_lib.MockTestCase):
 
     self.assertEquals(disk.DiskUsage(partition), mock_diskusage)
 
-  def testDiskByExisting(self):
-    """Test diskby information collection when a record exists."""
+  def testBlockDevicesExisting(self):
+    """Test block device information collection when a record exists."""
     disk = self._CreateDisk(1)
 
     device = '/dev/sda1'
-    dataname = '%s:%s' % (systeminfo.RESOURCENAME_DISKBY, device)
+    dataname = '%s:%s' % (systeminfo.RESOURCENAME_BLOCKDEVICE, device)
     data = 'testvalue'
 
     disk.Update(dataname, data)
 
-    self.assertEquals(disk.DiskBy(device), data)
+    self.assertEquals(disk.BlockDevices(device), data)
 
-  def testDiskBy(self):
-    """Test diskby info when there is no record or a record is stale."""
+  def testBlockDevice(self):
+    """Test block device info when there is no record or a record is stale."""
     disk = self._CreateDisk(1)
 
     mock_device = '/dev/sda1'
+    mock_size = 12345678987654321
     mock_ids = ['ata-ST1000DM003-1ER162_Z4Y3WQDB-part1']
     mock_labels = ['BOOT-PARTITION']
+    mock_lsblk = 'NAME="sda1" RM="0" TYPE="part" SIZE="%s"' % mock_size
+
+    self.StartPatcher(mock.patch('chromite.lib.osutils.ResolveSymlink'))
+    osutils.ResolveSymlink.return_value = '/dev/sda1'
 
-    # Mock the calls to cros_buil_lib.RunCommand.
     with cros_build_lib_unittest.RunCommandMock() as rc_mock:
       rc_mock.AddCmdResult(
           partial_mock.In(systeminfo.SYSTEMFILE_DEV_DISKBY['ids']),
@@ -372,11 +377,14 @@ class DiskTest(cros_test_lib.MockTestCase):
       rc_mock.AddCmdResult(
           partial_mock.In(systeminfo.SYSTEMFILE_DEV_DISKBY['labels']),
           output='\n'.join(mock_labels))
+      rc_mock.AddCmdResult(partial_mock.In('lsblk'), output=mock_lsblk)
 
-      mock_diskby = systeminfo.RESOURCE_DISKBY(mock_device, mock_ids,
-                                               mock_labels)
+      mock_blockdevice = [systeminfo.RESOURCE_BLOCKDEVICE(mock_device,
+                                                          mock_size,
+                                                          mock_ids,
+                                                          mock_labels)]
 
-      self.assertEquals(disk.DiskBy(mock_device), mock_diskby)
+      self.assertEquals(disk.BlockDevices(mock_device), mock_blockdevice)
 
 
 class Cpu(cros_test_lib.MockTestCase):