sweetberry: print MEAN in json format

This CL adds the flag --save_stats_json which prints the MEAN
of sweetberry readings into json format. Corresponding unit
test has been added too. Printing just the MEAN in json format
makes it easier to read it back in power_telemetry_logger,
which uploads the stats to the dashboard.

This CL is part of the effort to start sweetberry measurements
while running power autotests.

BRANCH=None
BUG=b:68956240
TEST=./powerlog.py -b xxx.board -c xxx.scenario \
--save_stats_json [directory to store it]
python -m  unittest -v stats_manager_unittest

Change-Id: I8274bcac21175f6c53184ced79b6ffe5e7d7a72a
Signed-off-by: Mengqi Guo <mqg@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/807636
Reviewed-by: Nick Sanders <nsanders@chromium.org>
Reviewed-by: Puthikorn Voravootivat <puthik@chromium.org>
Reviewed-by: Todd Broch <tbroch@chromium.org>

4 files changed, 75 insertions, 13 deletions

diff --git a/extra/usb_power/board.README b/extra/usb_power/board.README
index fbf9f527c5..b898e16540 100644
--- a/extra/usb_power/board.README
+++ b/extra/usb_power/board.README
@@ -114,3 +114,19 @@ under the directory which powerlog.py is in:
 <directory of powerlog.py>/sweetberry<timestemp>/raw_data/
 
 If --save_raw_data flag is not set, raw data will not be saved.
+
+Example 3:
+./powerlog.py -b board/eve_dvt2_loc/eve_dvt2_loc.board -c board/eve_dvt2_loc/eve_dvt2_loc.scenario --save_stats_json [<directory>]
+
+If <directory> is specified, this will save MEANS in json as:
+<directory>/sweetberry<timestemp>/summary.json
+If <directory> does not exist, it will be created.
+
+If <directory> is not specified but the flag is set, this will save MEANS in
+json under the directory which powerlog.py is in:
+<directory of powerlog.py>/sweetberry<timestemp>/summary.json
+
+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.
diff --git a/extra/usb_power/powerlog.py b/extra/usb_power/powerlog.py
index 5e8ed42c79..eaa2ffd18d 100755
--- a/extra/usb_power/powerlog.py
+++ b/extra/usb_power/powerlog.py
@@ -491,7 +491,8 @@ class powerlog(object):
 
   def __init__(self, brdfile, cfgfile, serial_a=None, serial_b=None,
                sync_date=False, use_ms=False, use_mW=False, print_stats=False,
-               stats_dir=None, print_raw_data=True, raw_data_dir=None):
+               stats_dir=None, stats_json_dir=None, print_raw_data=True,
+               raw_data_dir=None):
     """Init the powerlog class and set the variables.
 
     Args:
@@ -505,6 +506,8 @@ class powerlog(object):
       print_stats: print statistics for sweetberry readings at the end.
       stats_dir: directory to save sweetberry readings statistics; if None then
                  do not save the statistics.
+      stats_json_dir: directory to save means of sweetberry readings in json
+                      format; if None then do not save the statistics.
       print_raw_data: print sweetberry readings raw data in real time, default
                       is to print.
       raw_data_dir: directory to save sweetberry readings raw data; if None then
@@ -516,6 +519,7 @@ class powerlog(object):
     self._use_mW = use_mW
     self._print_stats = print_stats
     self._stats_dir = stats_dir
+    self._stats_json_dir = stats_json_dir
     self._print_raw_data = print_raw_data
     self._raw_data_dir = raw_data_dir
 
@@ -648,6 +652,9 @@ class powerlog(object):
       if self._stats_dir:
         stats_dir = os.path.join(self._stats_dir, save_dir)
         self._data.SaveSummary(stats_dir)
+      if self._stats_json_dir:
+        stats_json_dir = os.path.join(self._stats_json_dir, save_dir)
+        self._data.SaveSummaryJSON(stats_json_dir)
       if self._raw_data_dir:
         raw_data_dir = os.path.join(self._raw_data_dir, save_dir)
         self._data.SaveRawData(raw_data_dir)
@@ -685,23 +692,31 @@ def main(argv=None):
   parser.add_argument('--save_stats', type=str, nargs='?',
       dest='stats_dir', metavar='STATS_DIR',
       const=os.path.dirname(os.path.abspath(__file__)), default=None,
-      help="Save statistics for sweetberry readings to %(metavar)s if \
-      %(metavar)s is specified, %(metavar)s will be created if it does not \
-      exist; if %(metavar)s is not specified but the flag is set, stats will \
-      be saved to where %(prog)s is located; if this flag is not set, then do \
-      not save stats")
+      help="Save statistics for sweetberry readings to %(metavar)s if "
+           "%(metavar)s is specified, %(metavar)s will be created if it does "
+           "not exist; if %(metavar)s is not specified but the flag is set, "
+           "stats will be saved to where %(prog)s is located; if this flag is "
+           "not set, then do not save stats")
+  parser.add_argument('--save_stats_json', type=str, nargs='?',
+      dest='stats_json_dir', metavar='STATS_JSON_DIR',
+      const=os.path.dirname(os.path.abspath(__file__)), default=None,
+      help="Save means for sweetberry readings in json to %(metavar)s if "
+           "%(metavar)s is specified, %(metavar)s will be created if it does "
+           "not exist; if %(metavar)s is not specified but the flag is set, "
+           "stats will be saved to where %(prog)s is located; if this flag is "
+           "not set, then do not save stats")
   parser.add_argument('--no_print_raw_data',
       dest='print_raw_data', default=True, action="store_false",
-      help="Not print raw sweetberry readings at real time, default is to \
-      print")
+      help="Not print raw sweetberry readings at real time, default is to "
+           "print")
   parser.add_argument('--save_raw_data', type=str, nargs='?',
       dest='raw_data_dir', metavar='RAW_DATA_DIR',
       const=os.path.dirname(os.path.abspath(__file__)), default=None,
-      help="Save raw data for sweetberry readings to %(metavar)s if \
-      %(metavar)s is specified, %(metavar)s will be created if it does not \
-      exist; if %(metavar)s is not specified but the flag is set, raw data \
-      will be saved to where %(prog)s is located; if this flag is not set, \
-      then do not save raw data")
+      help="Save raw data for sweetberry readings to %(metavar)s if "
+           "%(metavar)s is specified, %(metavar)s will be created if it does "
+           "not exist; if %(metavar)s is not specified but the flag is set, "
+           "raw data will be saved to where %(prog)s is located; if this flag "
+           "is not set, then do not save raw data")
   parser.add_argument('-v', '--verbose', default=False,
       help="Very chatty printout", action="store_true")
 
@@ -727,6 +742,7 @@ def main(argv=None):
   use_mW = args.mW
   print_stats = args.print_stats
   stats_dir = args.stats_dir
+  stats_json_dir = args.stats_json_dir
   print_raw_data = args.print_raw_data
   raw_data_dir = args.raw_data_dir
 
@@ -740,6 +756,7 @@ def main(argv=None):
   powerlogger = powerlog(brdfile, cfgfile, serial_a=serial_a, serial_b=serial_b,
       sync_date=sync_date, use_ms=use_ms, use_mW=use_mW,
       print_stats=print_stats, stats_dir=stats_dir,
+      stats_json_dir=stats_json_dir,
       print_raw_data=print_raw_data,raw_data_dir=raw_data_dir)
 
   # Start logging.
diff --git a/extra/usb_power/stats_manager.py b/extra/usb_power/stats_manager.py
index 02f984f097..8e25c74a42 100644
--- a/extra/usb_power/stats_manager.py
+++ b/extra/usb_power/stats_manager.py
@@ -6,6 +6,7 @@
 
 from __future__ import print_function
 import collections
+import json
 import numpy
 import os
 
@@ -114,6 +115,24 @@ class StatsManager(object):
     with open(fname, 'w') as f:
       f.write(summary_str)
 
+  def SaveSummaryJSON(self, directory, fname='summary.json'):
+    """Save summary (only MEAN) into a JSON file.
+
+    Args:
+      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)
+    }
+    if not os.path.exists(directory):
+      os.makedirs(directory)
+    fname = os.path.join(directory, fname)
+    with open(fname, 'w') as f:
+      json.dump(data, f)
+
   def GetRawData(self):
     """Getter for all raw_data."""
     return self._data
diff --git a/extra/usb_power/stats_manager_unittest.py b/extra/usb_power/stats_manager_unittest.py
index 9b86b15ad4..7537368996 100644
--- a/extra/usb_power/stats_manager_unittest.py
+++ b/extra/usb_power/stats_manager_unittest.py
@@ -5,6 +5,7 @@
 """Unit tests for StatsManager."""
 
 from __future__ import print_function
+import json
 import os
 import shutil
 import tempfile
@@ -82,6 +83,15 @@ class TestStatsManager(unittest.TestCase):
           '@@      B      3       2.50    0.82       3.50      1.50\n',
           f.readline())
 
+  def test_SaveSummaryJSON(self):
+    fname = 'unittest_summary.json'
+    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'])
+
 
 if __name__ == '__main__':
   unittest.main()