sweetberry: fix stats_manager and refactor

This CL updates stats_manager to match the new functionalities
in powerlog.py and refactors powerlog.py to more easily find
config files and print timestamps in seconds since epoch.

The unit test for stats_manager is also updated accordingly.

BUG=b:72973433
BRANCH=None
TEST=powerlog -b nami_rev0_loc.board -c nami_rev0_loc.scenario \
--print_stats --save_stats /tmp --save_stats_json /tmp \
--save_raw_data /tmp --mW
and looking at the printed data
python -m unittest stats_manager_unittest
CQ-DEPEND=CL:1003522

Change-Id: Ic6e4aadfcd3ad245572788094ee3d3a30106044c
Signed-off-by: Mengqi Guo <mqg@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/1002546
Reviewed-by: Todd Broch <tbroch@chromium.org>

4 files changed, 149 insertions, 41 deletions

diff --git a/extra/usb_power/board.README b/extra/usb_power/board.README
index 41c30c8e38..fc9b055489 100644
--- a/extra/usb_power/board.README
+++ b/extra/usb_power/board.README
@@ -1,7 +1,7 @@
 Sweetberry USB power monitoring
 
 This tool allows high speed monitoring of power rails via a special USB
-endpoint. Currently this is implemented for the sweetberry baord.
+endpoint. Currently this is implemented for the sweetberry board.
 
 To use on a board, you'll need two config files, one describing the board,
 a ".board" file, and one describing the particular rails you want to
@@ -10,6 +10,8 @@ monitor in this session, a ".scenario" file.
 
 Converting from servo_ina configs:
 
+Method 1 -
+
 Many configs can be found for the servo_ina_board in hdctools/servo/data/.
 Sweetberry is plug compatible with servo_ina headers, and config files
 can be converted with the following tool:
@@ -19,6 +21,20 @@ can be converted with the following tool:
 This will produce kevin_r0_loc.board and kevin_r0_loc.scenario which
 can be used with powerlog.py.
 
+Method 2 (preferred) -
+
+If you are using powerlog.py within the chroot, copy kevin_r0_loc.py to
+src/third_party/hdctools/servo/data, then add line to file:
+config_type = 'sweetberry'
+and run command in terminal:
+sudo emerge hdctools
+The command will install the corresponding .board and .scenario file in the
+chroot. To use powerlog.py use the command:
+./powerlog.py -b kevin_r0_loc.board -c kevin_r0_loc.scenario
+There is no need to specify the absolute path to the .board and .scenario file,
+once they are installed into the chroot. If there is any changes to
+kevin_r0_loc.py, you need to emerge hdctools again.
+
 
 Board files:
 
@@ -137,3 +153,15 @@ If --save_stats flag is not set, stats will not be saved.
 
 --save_stats_json is designed for power_telemetry_logger for easy reading and
 writing.
+
+
+Making developer changes to powerlog.py:
+
+powerlog.py is installed in chroot, and the developer can import powerlog or use
+powerlog directly anywhere within chroot. Anytime the developer makes a change
+to powerlog.py, the developer needs to re-install powerlog.py so that anything
+that imports powerlog does not break. The following is how the developer
+installs powerlog.py during development.
+Run command in the terminal:
+cros_workon --host start ec-devutils # just the first time
+sudo emerge ec-devutils # everytime powerlog gets changed
diff --git a/extra/usb_power/powerlog.py b/extra/usb_power/powerlog.py
index 6128ca2027..00a508691a 100755
--- a/extra/usb_power/powerlog.py
+++ b/extra/usb_power/powerlog.py
@@ -10,7 +10,7 @@
 from __future__ import print_function
 import argparse
 import array
-import datetime
+from distutils import sysconfig
 import json
 import os
 import struct
@@ -23,6 +23,9 @@ import usb
 
 from stats_manager import StatsManager
 
+LIB_DIR = os.path.join(sysconfig.get_python_lib(standard_lib=False), 'servo',
+                       'data')
+
 # This can be overridden by -v.
 debug = False
 def debuglog(msg):
@@ -33,6 +36,40 @@ def logoutput(msg):
   print(msg)
   sys.stdout.flush()
 
+def process_filename(filename):
+  """Find the file path from the filename.
+
+  If filename is already the complete path, return that directly. If filename is
+  just the short name, look for the file in the current working directory, in
+  the directory of the current .py file, and then in the directory installed by
+  hdctools. If the file is found, return the complete path of the file.
+
+  Args:
+      filename: complete file path or short file name.
+
+  Returns:
+      a complete file path.
+
+  Raises:
+      IOError if filename does not exist.
+  """
+  # 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
+  raise IOError('No such file or directory: \'%s\'' % filename)
+
 
 class Spower(object):
   """Power class to access devices on the bus.
@@ -51,6 +88,10 @@ class Spower(object):
   INA_BUSV      = 2
   INA_CURRENT   = 3
   INA_SHUNTV    = 4
+  # INA_SUFFIX is used to differentiate multiple ina types for the same power
+  # rail. No suffix for when ina type is 0 (non-existent) and when ina type is 1
+  # (power, no suffix for backward compatibility).
+  INA_SUFFIX      = ['', '', '_busv', '_cur', '_shuntv']
 
   # usb power commands
   CMD_RESET   = 0x0000
@@ -503,7 +544,7 @@ class Spower(object):
     Args:
       brdfile:	Filename of a json file decribing the INA wiring of this board.
     """
-    with open(brdfile) as data_file:
+    with open(process_filename(brdfile)) as data_file:
         data = json.load(data_file)
 
     #TODO: validate this.
@@ -519,7 +560,8 @@ class powerlog(object):
     obj = powerlog()
 
   Instance Variables:
-    _data: records sweetberries readings and calculates statistics.
+    _data: a StatsManager object that records sweetberry readings and calculates
+           statistics.
     _pwr[]: Spower objects for individual sweetberries.
   """
 
@@ -564,7 +606,7 @@ class powerlog(object):
     if serial_b:
       self._pwr['B'] = Spower('B', serialname=serial_b)
 
-    with open(cfgfile) as data_file:
+    with open(process_filename(cfgfile)) as data_file:
       names = json.load(data_file)
     self._names = self.process_scenario(names)
 
@@ -648,22 +690,24 @@ class powerlog(object):
       integration_us = integration_us_new
 
     # CSV header
+    title = "ts:%dus" % integration_us
+    for name_tuple in self._names:
+      name, ina_type = name_tuple
+
+      if ina_type == Spower.INA_POWER:
+        unit = "mW" if self._use_mW else "uW"
+      elif ina_type == Spower.INA_BUSV:
+        unit = "mV"
+      elif ina_type == Spower.INA_CURRENT:
+        unit = "uA"
+      elif ina_type == Spower.INA_SHUNTV:
+        unit = "uV"
+
+      title += ", %s %s" % (name, unit)
+      name_type = name + Spower.INA_SUFFIX[ina_type]
+      self._data.SetUnit(name_type, unit)
+    title += ", status"
     if self._print_raw_data:
-      title = "ts:%dus" % integration_us
-      for name_tuple in self._names:
-        name, ina_type = name_tuple
-
-        if ina_type == Spower.INA_POWER:
-          unit = "mW" if self._use_mW else "uW"
-        elif ina_type == Spower.INA_BUSV:
-          unit = "mV"
-        elif ina_type == Spower.INA_CURRENT:
-          unit = "uA"
-        elif ina_type == Spower.INA_SHUNTV:
-          unit = "uV"
-
-        title += ", %s %s" % (name, unit)
-      title += ", status"
       logoutput(title)
 
     forever = False
@@ -704,7 +748,8 @@ class powerlog(object):
                     name[1]==Spower.INA_POWER) else 1
                 value = aggregate_record[name] * multiplier
                 csv += ", %.2f" % value
-                self._data.AddValue(name, value)
+                name_type = name[0] + Spower.INA_SUFFIX[name[1]]
+                self._data.AddValue(name_type, value)
               else:
                 csv += ", "
             csv += ", %d" % aggregate_record["status"]
@@ -724,7 +769,7 @@ class powerlog(object):
       self._data.CalculateStats()
       if self._print_stats:
         self._data.PrintSummary()
-      save_dir = datetime.datetime.now().strftime('sweetberry%Y%m%d%H%M%S.%f')
+      save_dir = 'sweetberry%s' % time.time()
       if self._stats_dir:
         stats_dir = os.path.join(self._stats_dir, save_dir)
         self._data.SaveSummary(stats_dir)
diff --git a/extra/usb_power/stats_manager.py b/extra/usb_power/stats_manager.py
index 8e25c74a42..3a85935a36 100644
--- a/extra/usb_power/stats_manager.py
+++ b/extra/usb_power/stats_manager.py
@@ -16,12 +16,22 @@ KEY_PREFIX = '__'
 # as timeline keys.
 NOSHOW_PREFIX = '!!'
 
+LONG_UNIT = {
+    'mW': 'milliwatt',
+    'uW': 'microwatt',
+    'mV': 'millivolt',
+    'uA': 'microamp',
+    'uV': 'microvolt'
+}
+
+
 class StatsManager(object):
   """Calculates statistics for several lists of data(float)."""
 
   def __init__(self):
     """Initialize infrastructure for data and their statistics."""
     self._data = collections.defaultdict(list)
+    self._unit = {}
     self._summary = {}
 
   def AddValue(self, domain, value):
@@ -39,6 +49,21 @@ class StatsManager(object):
     print('Warning: value %s for domain %s is not a number, thus ignored.' %
           (value, domain))
 
+  def SetUnit(self, domain, unit):
+    """Set the unit for a domain.
+
+    There can be only one unit for each domain. Setting unit twice will
+    overwrite the original unit.
+
+    Args:
+      domain: the domain name.
+      unit: unit of the domain.
+    """
+    if domain in self._unit:
+      print('Warning: overwriting the unit of %s, old unit is %s, new unit is '
+            '%s.' % (domain, self._unit[domain], unit))
+    self._unit[domain] = unit
+
   def CalculateStats(self):
     """Calculate stats for all domain-data pairs.
 
@@ -48,11 +73,11 @@ class StatsManager(object):
     for domain, data in self._data.iteritems():
       data_np = numpy.array(data)
       self._summary[domain] = {
-          'mean'     : data_np.mean(),
-          'min'      : data_np.min(),
-          'max'      : data_np.max(),
-          'stddev'   : data_np.std(),
-          'count'    : data_np.size,
+          'mean': data_np.mean(),
+          'min': data_np.min(),
+          'max': data_np.max(),
+          'stddev': data_np.std(),
+          'count': data_np.size,
       }
 
   def _SummaryToString(self, prefix=STATS_PREFIX):
@@ -67,7 +92,9 @@ class StatsManager(object):
       if domain.startswith(NOSHOW_PREFIX):
         continue
       stats = self._summary[domain]
-      row = [domain.lstrip(KEY_PREFIX)]
+      unit = self._unit[domain]
+      domain_unit = domain.lstrip(KEY_PREFIX) + '_' + unit
+      row = [domain_unit]
       row.append(str(stats['count']))
       for entry in headers[2:]:
         row.append('%.2f' % stats[entry.lower()])
@@ -122,11 +149,13 @@ class StatsManager(object):
       directory: directory to save the JSON summary in.
       fname: filename to save summary under.
     """
-    data = {
-        domain: self._summary[domain]['mean']
-        for domain in sorted(self._summary.keys())
-        if not domain.startswith(NOSHOW_PREFIX)
-    }
+    data = {}
+    for domain in self._summary:
+      if domain.startswith(NOSHOW_PREFIX):
+        continue
+      unit = LONG_UNIT.get(self._unit[domain], self._unit[domain])
+      data_entry = {'mean': self._summary[domain]['mean'], 'unit': unit}
+      data[domain] = data_entry
     if not os.path.exists(directory):
       os.makedirs(directory)
     fname = os.path.join(directory, fname)
@@ -150,7 +179,7 @@ class StatsManager(object):
     if not os.path.exists(dirname):
       os.makedirs(dirname)
     for domain, data in self._data.iteritems():
-      fname = domain + '.txt'
+      fname = domain + '_' + self._unit[domain] + '.txt'
       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 7537368996..bf0861ff2d 100644
--- a/extra/usb_power/stats_manager_unittest.py
+++ b/extra/usb_power/stats_manager_unittest.py
@@ -13,6 +13,7 @@ import unittest
 
 from stats_manager import StatsManager
 
+
 class TestStatsManager(unittest.TestCase):
   """Test to verify StatsManager methods work as expected.
 
@@ -27,9 +28,12 @@ class TestStatsManager(unittest.TestCase):
     self.data.AddValue('A', 99999.5)
     self.data.AddValue('A', 100000.5)
     self.data.AddValue('A', 'ERROR')
+    self.data.SetUnit('A', 'uW')
+    self.data.SetUnit('A', 'mW')
     self.data.AddValue('B', 1.5)
     self.data.AddValue('B', 2.5)
     self.data.AddValue('B', 3.5)
+    self.data.SetUnit('B', 'mV')
     self.data.CalculateStats()
 
   def tearDown(self):
@@ -58,8 +62,8 @@ class TestStatsManager(unittest.TestCase):
     dirname = 'unittest_raw_data'
     self.data.SaveRawData(self.tempdir, dirname)
     dirname = os.path.join(self.tempdir, dirname)
-    fileA = os.path.join(dirname, 'A.txt')
-    fileB = os.path.join(dirname, 'B.txt')
+    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())
@@ -77,10 +81,10 @@ class TestStatsManager(unittest.TestCase):
           '@@   NAME  COUNT       MEAN  STDDEV        MAX       MIN\n',
           f.readline())
       self.assertEqual(
-          '@@      A      2  100000.00    0.50  100000.50  99999.50\n',
+          '@@   A_mW      2  100000.00    0.50  100000.50  99999.50\n',
           f.readline())
       self.assertEqual(
-          '@@      B      3       2.50    0.82       3.50      1.50\n',
+          '@@   B_mV      3       2.50    0.82       3.50      1.50\n',
           f.readline())
 
   def test_SaveSummaryJSON(self):
@@ -88,9 +92,11 @@ class TestStatsManager(unittest.TestCase):
     self.data.SaveSummaryJSON(self.tempdir, fname)
     fname = os.path.join(self.tempdir, fname)
     with open(fname, 'r') as f:
-      mean_json = json.load(f)
-      self.assertAlmostEqual(100000.0, mean_json['A'])
-      self.assertAlmostEqual(2.5, mean_json['B'])
+      summary = json.load(f)
+      self.assertAlmostEqual(100000.0, summary['A']['mean'])
+      self.assertEqual('milliwatt', summary['A']['unit'])
+      self.assertAlmostEqual(2.5, summary['B']['mean'])
+      self.assertEqual('millivolt', summary['B']['unit'])
 
 
 if __name__ == '__main__':