stats_manager: accept_nan support

This CL introduces a flag to StatsManager that allows for 'NaN' values
to be recorded inside StatsManager. The motivation here is that if a
sample fails to record it might be more desirable to record a 'NaN'
than to just skip the sample, to keep timelines correct, and to not hide
errors in the test-run.
Also adds necessary tests for that behavior.

BRANCH=None
BUG=chromium:806146, chromium:760267
TEST=unit tests still pass

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

2 files changed, 45 insertions, 16 deletions

diff --git a/extra/usb_power/stats_manager.py b/extra/usb_power/stats_manager.py
index 66101ab92e..66dd915a86 100644
--- a/extra/usb_power/stats_manager.py
+++ b/extra/usb_power/stats_manager.py
@@ -9,6 +9,7 @@ from __future__ import print_function
 import collections
 import json
 import logging
+import math
 import os
 
 import numpy
@@ -76,7 +77,8 @@ class StatsManager(object):
   """
 
   # pylint: disable=W0102
-  def __init__(self, smid='', title='', order=[], hide_domains=[]):
+  def __init__(self, smid='', title='', order=[], hide_domains=[],
+               accept_nan=True):
     """Initialize infrastructure for data and their statistics."""
     self._title = title
     self._data = collections.defaultdict(list)
@@ -85,6 +87,7 @@ class StatsManager(object):
     self._order = order
     self._hide_domains = hide_domains
     self._summary = {}
+    self._accept_nan = accept_nan
     self._logger = logging.getLogger('StatsManager')
 
   def AddSample(self, domain, sample):
@@ -93,14 +96,20 @@ class StatsManager(object):
     Args:
       domain: the domain name for the sample.
       sample: one time sample for domain, expect type float.
+
+    Raises:
+      StatsManagerError: if trying to add NaN and |_accept_nan| is false
     """
-    if isinstance(sample, int):
+    try:
       sample = float(sample)
-    if isinstance(sample, float):
-      self._data[domain].append(sample)
-      return
-    self._logger.warn('sample %s for domain %s is not a number, thus ignored.',
-                      sample, domain)
+    except ValueError:
+      # if we don't accept nan this will be caught below
+      self._logger.debug('sample %s for domain %s is not a number. Making NaN',
+                         sample, domain)
+      sample = float('NaN')
+    if not self._accept_nan and math.isnan(sample):
+      raise StatsManagerError('accept_nan is false. Cannot add NaN sample.')
+    self._data[domain].append(sample)
 
   def SetUnit(self, domain, unit):
     """Set the unit for a domain.
@@ -126,10 +135,10 @@ 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(),
+          'mean': numpy.nanmean(data_np),
+          'min': numpy.nanmin(data_np),
+          'max': numpy.nanmax(data_np),
+          'stddev': numpy.nanstd(data_np),
           'count': data_np.size,
       }
 
diff --git a/extra/usb_power/stats_manager_unittest.py b/extra/usb_power/stats_manager_unittest.py
index e2dda06f91..ed0b424944 100644
--- a/extra/usb_power/stats_manager_unittest.py
+++ b/extra/usb_power/stats_manager_unittest.py
@@ -26,7 +26,6 @@ class TestStatsManager(unittest.TestCase):
     """Create a populated & processed StatsManager to test data retrieval."""
     self.data.AddSample('A', 99999.5)
     self.data.AddSample('A', 100000.5)
-    self.data.AddSample('A', 'ERROR')
     self.data.SetUnit('A', 'uW')
     self.data.SetUnit('A', 'mW')
     self.data.AddSample('B', 1.5)
@@ -57,14 +56,35 @@ class TestStatsManager(unittest.TestCase):
     summary = self.data.GetSummary()
     self.assertEqual(1, summary['Test']['count'])
 
-  def test_AddSampleNoFloat(self):
-    """Adding a non number gets ignored and doesn't raise an exception."""
-    self.data.AddSample('Test', 17)
+  def test_AddSampleNoFloatAcceptNaN(self):
+    """Adding a non-number adds 'NaN' and doesn't raise an exception."""
+    self.data.AddSample('Test', 10)
+    self.data.AddSample('Test', 20)
+    # adding a fake NaN: one that gets converted into NaN internally
     self.data.AddSample('Test', 'fiesta')
+    # adding a real NaN
+    self.data.AddSample('Test', float('NaN'))
     self.data.SetUnit('Test', 'test')
     self.data.CalculateStats()
     summary = self.data.GetSummary()
-    self.assertEqual(1, summary['Test']['count'])
+    # assert that 'NaN' as added.
+    self.assertEqual(4, summary['Test']['count'])
+    # assert that mean, min, and max calculatings ignore the 'NaN'
+    self.assertEqual(10, summary['Test']['min'])
+    self.assertEqual(20, summary['Test']['max'])
+    self.assertEqual(15, summary['Test']['mean'])
+
+  def test_AddSampleNoFloatNotAcceptNaN(self):
+    """Adding a non-number raises a StatsManagerError if accept_nan is False."""
+    self.data = stats_manager.StatsManager(accept_nan=False)
+    with self.assertRaisesRegexp(stats_manager.StatsManagerError,
+                                 'accept_nan is false. Cannot add NaN sample.'):
+      # adding a fake NaN: one that gets converted into NaN internally
+      self.data.AddSample('Test', 'fiesta')
+    with self.assertRaisesRegexp(stats_manager.StatsManagerError,
+                                 'accept_nan is false. Cannot add NaN sample.'):
+      # adding a real NaN
+      self.data.AddSample('Test', float('NaN'))
 
   def test_AddSampleNoUnit(self):
     """Not adding a unit does not cause an exception on CalculateStats()."""