stats_manager: prepare StatsManager to be a utility used in hdctools

This is the first CL in a series of CLs to start using StatsManager in
servo/hdctools (package depends on ec-devutils, as in this package).
This CL:
- beefs up StatsManager to handle unavailable units more gracefully
- adds a few more tests to stats_manager_unittest.py
- adds some minor unit testing for powerlog's file retrieval logic

BRANCH=None
BUG=chromium:760267
TEST=manual testing, unit tests still pass, powerlog still works

Change-Id: Ifcdfcc482008484fbc21326c6f087ebf466c3e74
Signed-off-by: Ruben Rodriguez Buchillon <coconutruben@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/1140025
Reviewed-by: Mengqi Guo <mqg@chromium.org>

4 files changed, 207 insertions, 34 deletions

diff --git a/extra/usb_power/powerlog.py b/extra/usb_power/powerlog.py
index 82950100f2..6d7a8bab26 100755
--- a/extra/usb_power/powerlog.py
+++ b/extra/usb_power/powerlog.py
@@ -23,9 +23,15 @@ import usb
 
 from stats_manager import StatsManager
 
+# Directory where hdctools installs configuration files into.
 LIB_DIR = os.path.join(sysconfig.get_python_lib(standard_lib=False), 'servo',
                        'data')
 
+# Potential config file locations: current working directory, the same directory
+# as powerlog.py file or LIB_DIR.
+CONFIG_LOCATIONS = [os.getcwd(), os.path.dirname(os.path.realpath(__file__)),
+                    LIB_DIR]
+
 # This can be overridden by -v.
 debug = False
 def debuglog(msg):
@@ -56,18 +62,11 @@ def process_filename(filename):
   # Check if filename is absolute path.
   if os.path.isabs(filename) and os.path.isfile(filename):
     return filename
-  # Check if filename is relative to current working directory.
-  cwd = os.path.join(os.getcwd(), filename)
-  if os.path.isfile(cwd):
-    return cwd
-  # Check if filename is relative to same directory as current .py file.
-  sd = os.path.join(os.path.dirname(os.path.realpath(__file__)), filename)
-  if os.path.isfile(sd):
-    return sd
-  # Check if file is installed by hdctools.
-  hdc = os.path.join(LIB_DIR, filename)
-  if os.path.isfile(hdc):
-    return hdc
+  # Check if filename is relative to a known config location.
+  for dirname in CONFIG_LOCATIONS:
+    file_at_dir = os.path.join(dirname, filename)
+    if os.path.isfile(file_at_dir):
+      return file_at_dir
   raise IOError('No such file or directory: \'%s\'' % filename)
 
 
diff --git a/extra/usb_power/powerlog_unittest.py b/extra/usb_power/powerlog_unittest.py
new file mode 100644
index 0000000000..7058c57aa7
--- /dev/null
+++ b/extra/usb_power/powerlog_unittest.py
@@ -0,0 +1,49 @@
+# Copyright 2018 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+"""Unit tests for powerlog."""
+
+import os
+import shutil
+import tempfile
+import unittest
+
+import powerlog
+
+class TestPowerlog(unittest.TestCase):
+  """Test to verify powerlog util methods work as expected."""
+
+  def setUp(self):
+    """Set up data and create a temporary directory to save data and stats."""
+    self.tempdir = tempfile.mkdtemp()
+    self.filename = 'testfile'
+    self.filepath = os.path.join(self.tempdir, self.filename)
+    with open(self.filepath, 'w') as f:
+      f.write('')
+
+  def tearDown(self):
+    """Delete the temporary directory and its content."""
+    shutil.rmtree(self.tempdir)
+
+  def test_ProcessFilenameAbsoluteFilePath(self):
+    """Absolute file path is returned unchanged."""
+    processed_fname = powerlog.process_filename(self.filepath)
+    self.assertEqual(self.filepath, processed_fname)
+
+  def test_ProcessFilenameRelativeFilePath(self):
+    """Finds relative file path inside a known config location."""
+    original = powerlog.CONFIG_LOCATIONS
+    powerlog.CONFIG_LOCATIONS = [self.tempdir]
+    processed_fname = powerlog.process_filename(self.filename)
+    try:
+      self.assertEqual(self.filepath, processed_fname)
+    finally:
+      powerlog.CONFIG_LOCATIONS = original
+
+  def test_ProcessFilenameInvalid(self):
+    """IOError is raised when file cannot be found by any of the four ways."""
+    with self.assertRaises(IOError):
+      powerlog.process_filename(self.filename)
+
+if __name__ == '__main__':
+  unittest.main()
diff --git a/extra/usb_power/stats_manager.py b/extra/usb_power/stats_manager.py
index 3a85935a36..a53f555832 100644
--- a/extra/usb_power/stats_manager.py
+++ b/extra/usb_power/stats_manager.py
@@ -5,18 +5,22 @@
 """Calculates statistics for lists of data and pretty print them."""
 
 from __future__ import print_function
+
 import collections
 import json
-import numpy
 import os
 
+import numpy
+
 STATS_PREFIX = '@@'
+# used to aid sorting of dict keys
 KEY_PREFIX = '__'
 # This prefix is used for keys that should not be shown in the summary tab, such
 # as timeline keys.
 NOSHOW_PREFIX = '!!'
 
 LONG_UNIT = {
+    '': 'N/A',
     'mW': 'milliwatt',
     'uW': 'microwatt',
     'mV': 'millivolt',
@@ -25,13 +29,47 @@ LONG_UNIT = {
 }
 
 
+class StatsManagerError(Exception):
+  """Errors in StatsManager class."""
+  pass
+
+
 class StatsManager(object):
-  """Calculates statistics for several lists of data(float)."""
+  """Calculates statistics for several lists of data(float).
+
+  Example usage:
+
+    >>> stats = StatsManager()
+    >>> stats.AddValue(TIME_KEY, 50.0)
+    >>> stats.AddValue(TIME_KEY, 25.0)
+    >>> stats.AddValue(TIME_KEY, 40.0)
+    >>> stats.AddValue(TIME_KEY, 10.0)
+    >>> stats.AddValue(TIME_KEY, 10.0)
+    >>> stats.AddValue('frobnicate', 11.5)
+    >>> stats.AddValue('frobnicate', 9.0)
+    >>> stats.AddValue('foobar', 11111.0)
+    >>> stats.AddValue('foobar', 22222.0)
+    >>> stats.CalculateStats()
+    >>> stats.PrintSummary()
+    @@            NAME  COUNT      MEAN   STDDEV       MAX       MIN
+    @@   sample_msecs      4     31.25    15.16     50.00     10.00
+    @@         foobar      2  16666.50  5555.50  22222.00  11111.00
+    @@     frobnicate      2     10.25     1.25     11.50      9.00
+
+  Attributes:
+    _data: dict of list of readings for each domain(key)
+    _unit: dict of unit for each domain(key)
+    _summary: dict of stats per domain (key): min, max, count, mean, stddev
+
+  Note:
+    _summary is empty until CalculateStats() is called, and is updated when
+    CalculateStats() is called.
+  """
 
   def __init__(self):
     """Initialize infrastructure for data and their statistics."""
     self._data = collections.defaultdict(list)
-    self._unit = {}
+    self._unit = collections.defaultdict(str)
     self._summary = {}
 
   def AddValue(self, domain, value):
@@ -83,8 +121,13 @@ class StatsManager(object):
   def _SummaryToString(self, prefix=STATS_PREFIX):
     """Format summary into a string, ready for pretty print.
 
+    See class description for format example.
+
     Args:
       prefix: start every row in summary string with prefix, for easier reading.
+
+    Returns:
+      formatted summary string.
     """
     headers = ('NAME', 'COUNT', 'MEAN', 'STDDEV', 'MAX', 'MIN')
     table = [headers]
@@ -92,9 +135,9 @@ class StatsManager(object):
       if domain.startswith(NOSHOW_PREFIX):
         continue
       stats = self._summary[domain]
-      unit = self._unit[domain]
-      domain_unit = domain.lstrip(KEY_PREFIX) + '_' + unit
-      row = [domain_unit]
+      if not domain.endswith(self._unit[domain]):
+        domain = '%s_%s' % (domain, self._unit[domain])
+      row = [domain.lstrip(KEY_PREFIX)]
       row.append(str(stats['count']))
       for entry in headers[2:]:
         row.append('%.2f' % stats[entry.lower()])
@@ -179,7 +222,9 @@ class StatsManager(object):
     if not os.path.exists(dirname):
       os.makedirs(dirname)
     for domain, data in self._data.iteritems():
-      fname = domain + '_' + self._unit[domain] + '.txt'
+      if not domain.endswith(self._unit[domain]):
+        domain = '%s_%s' % (domain, self._unit[domain])
+      fname = '%s.txt' % domain
       fname = os.path.join(dirname, fname)
       with open(fname, 'w') as f:
         f.write('\n'.join('%.2f' % value for value in data) + '\n')
diff --git a/extra/usb_power/stats_manager_unittest.py b/extra/usb_power/stats_manager_unittest.py
index bf0861ff2d..362db3fa77 100644
--- a/extra/usb_power/stats_manager_unittest.py
+++ b/extra/usb_power/stats_manager_unittest.py
@@ -11,8 +11,7 @@ import shutil
 import tempfile
 import unittest
 
-from stats_manager import StatsManager
-
+import stats_manager
 
 class TestStatsManager(unittest.TestCase):
   """Test to verify StatsManager methods work as expected.
@@ -21,10 +20,8 @@ class TestStatsManager(unittest.TestCase):
   them in expected format.
   """
 
-  def setUp(self):
-    """Set up data and create a temporary directory to save data and stats."""
-    self.tempdir = tempfile.mkdtemp()
-    self.data = StatsManager()
+  def _populate_dummy_stats(self):
+    """Create a populated & processed StatsManager to test data retrieval."""
     self.data.AddValue('A', 99999.5)
     self.data.AddValue('A', 100000.5)
     self.data.AddValue('A', 'ERROR')
@@ -36,16 +33,71 @@ class TestStatsManager(unittest.TestCase):
     self.data.SetUnit('B', 'mV')
     self.data.CalculateStats()
 
+  def _populate_dummy_stats_no_unit(self):
+    self.data.AddValue('B', 1000)
+    self.data.AddValue('A', 200)
+    self.data.SetUnit('A', 'blue')
+
+  def setUp(self):
+    """Set up StatsManager and create a temporary directory for test."""
+    self.tempdir = tempfile.mkdtemp()
+    self.data = stats_manager.StatsManager()
+
   def tearDown(self):
     """Delete the temporary directory and its content."""
     shutil.rmtree(self.tempdir)
 
+  def test_AddValue(self):
+    """Adding a value successfully adds a value."""
+    self.data.AddValue('Test', 1000)
+    self.data.SetUnit('Test', 'test')
+    self.data.CalculateStats()
+    summary = self.data.GetSummary()
+    self.assertEqual(1, summary['Test']['count'])
+
+  def test_AddValueNoFloat(self):
+    """Adding a non number gets ignored and doesn't raise an exception."""
+    self.data.AddValue('Test', 17)
+    self.data.AddValue('Test', 'fiesta')
+    self.data.SetUnit('Test', 'test')
+    self.data.CalculateStats()
+    summary = self.data.GetSummary()
+    self.assertEqual(1, summary['Test']['count'])
+
+  def test_AddValueNoUnit(self):
+    """Not adding a unit does not cause an exception on CalculateStats()."""
+    self.data.AddValue('Test', 17)
+    self.data.CalculateStats()
+    summary = self.data.GetSummary()
+    self.assertEqual(1, summary['Test']['count'])
+
+  def test_UnitSuffix(self):
+    """Unit gets appended as a suffix in the displayed summary."""
+    self.data.AddValue('test', 250)
+    self.data.SetUnit('test', 'mw')
+    self.data.CalculateStats()
+    summary_str = self.data._SummaryToString()
+    self.assertIn('test_mw', summary_str)
+
+  def test_DoubleUnitSuffix(self):
+    """If domain already ends in unit, verify that unit doesn't get appended."""
+    self.data.AddValue('test_mw', 250)
+    self.data.SetUnit('test_mw', 'mw')
+    self.data.CalculateStats()
+    summary_str = self.data._SummaryToString()
+    self.assertIn('test_mw', summary_str)
+    self.assertNotIn('test_mw_mw', summary_str)
+
   def test_GetRawData(self):
+    """GetRawData returns exact same data as fed in."""
+    self._populate_dummy_stats()
     raw_data = self.data.GetRawData()
     self.assertListEqual([99999.5, 100000.5], raw_data['A'])
     self.assertListEqual([1.5, 2.5, 3.5], raw_data['B'])
 
   def test_GetSummary(self):
+    """GetSummary returns expected stats about the data fed in."""
+    self._populate_dummy_stats()
     summary = self.data.GetSummary()
     self.assertEqual(2, summary['A']['count'])
     self.assertAlmostEqual(100000.5, summary['A']['max'])
@@ -59,20 +111,34 @@ class TestStatsManager(unittest.TestCase):
     self.assertAlmostEqual(2.5, summary['B']['mean'])
 
   def test_SaveRawData(self):
+    """SaveRawData stores same data as fed in."""
+    self._populate_dummy_stats()
     dirname = 'unittest_raw_data'
     self.data.SaveRawData(self.tempdir, dirname)
     dirname = os.path.join(self.tempdir, dirname)
-    fileA = os.path.join(dirname, 'A_mW.txt')
-    fileB = os.path.join(dirname, 'B_mV.txt')
-    with open(fileA, 'r') as fA:
-      self.assertEqual('99999.50', fA.readline().strip())
-      self.assertEqual('100000.50', fA.readline().strip())
-    with open(fileB, 'r') as fB:
-      self.assertEqual('1.50', fB.readline().strip())
-      self.assertEqual('2.50', fB.readline().strip())
-      self.assertEqual('3.50', fB.readline().strip())
+    file_a = os.path.join(dirname, 'A_mW.txt')
+    file_b = os.path.join(dirname, 'B_mV.txt')
+    with open(file_a, 'r') as f_a:
+      self.assertEqual('99999.50', f_a.readline().strip())
+      self.assertEqual('100000.50', f_a.readline().strip())
+    with open(file_b, 'r') as f_b:
+      self.assertEqual('1.50', f_b.readline().strip())
+      self.assertEqual('2.50', f_b.readline().strip())
+      self.assertEqual('3.50', f_b.readline().strip())
+
+  def test_SaveRawDataNoUnit(self):
+    """SaveRawData appends no unit suffix if the unit is not specified."""
+    self._populate_dummy_stats_no_unit()
+    self.data.CalculateStats()
+    outdir = 'unittest_raw_data'
+    self.data.SaveRawData(self.tempdir, outdir)
+    files = os.listdir(os.path.join(self.tempdir, outdir))
+    # Verify nothing gets appended to domain for filename if no unit exists.
+    self.assertIn('B.txt', files)
 
   def test_SaveSummary(self):
+    """SaveSummary properly dumps the summary into a file."""
+    self._populate_dummy_stats()
     fname = 'unittest_summary.txt'
     self.data.SaveSummary(self.tempdir, fname)
     fname = os.path.join(self.tempdir, fname)
@@ -88,6 +154,8 @@ class TestStatsManager(unittest.TestCase):
           f.readline())
 
   def test_SaveSummaryJSON(self):
+    """SaveSummaryJSON saves the added data properly in JSON format."""
+    self._populate_dummy_stats()
     fname = 'unittest_summary.json'
     self.data.SaveSummaryJSON(self.tempdir, fname)
     fname = os.path.join(self.tempdir, fname)
@@ -98,6 +166,18 @@ class TestStatsManager(unittest.TestCase):
       self.assertAlmostEqual(2.5, summary['B']['mean'])
       self.assertEqual('millivolt', summary['B']['unit'])
 
+  def test_SaveSummaryJSONNoUnit(self):
+    """SaveSummaryJSON marks unknown units properly as N/A."""
+    self._populate_dummy_stats_no_unit()
+    self.data.CalculateStats()
+    fname = 'unittest_summary.json'
+    self.data.SaveSummaryJSON(self.tempdir, fname)
+    fname = os.path.join(self.tempdir, fname)
+    with open(fname, 'r') as f:
+      summary = json.load(f)
+      self.assertEqual('blue', summary['A']['unit'])
+      # if no unit is specified, JSON should save 'N/A' as the unit.
+      self.assertEqual('N/A', summary['B']['unit'])
 
 if __name__ == '__main__':
   unittest.main()