stats_manager: StatsManager returns file save locations

Some test code currently duplicates knowledge by telling StatsManager
where to save a file and then having to figure out where it stored the
files again. This CL has StatsManager return the save locations to
simplfy code. This potentially allows for informative logging on where
information was stored.

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

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

2 files changed, 48 insertions, 19 deletions

diff --git a/extra/usb_power/stats_manager.py b/extra/usb_power/stats_manager.py
index 92d15098bc..cc663ecce0 100644
--- a/extra/usb_power/stats_manager.py
+++ b/extra/usb_power/stats_manager.py
@@ -174,6 +174,9 @@ class StatsManager(object):
       directory: directory to save the summary in.
       fname: filename to save summary under.
       prefix: start every row in summary string with prefix, for easier reading.
+
+    Returns:
+      full path of summary save location
     """
     summary_str = self.SummaryToString(prefix=prefix) + '\n'
 
@@ -182,6 +185,7 @@ class StatsManager(object):
     fname = os.path.join(directory, fname)
     with open(fname, 'w') as f:
       f.write(summary_str)
+    return fname
 
   def SaveSummaryJSON(self, directory, fname='summary.json'):
     """Save summary (only MEAN) into a JSON file.
@@ -189,6 +193,9 @@ class StatsManager(object):
     Args:
       directory: directory to save the JSON summary in.
       fname: filename to save summary under.
+
+    Returns:
+      full path of summary save location
     """
     data = {}
     for domain in self._summary:
@@ -200,6 +207,7 @@ class StatsManager(object):
     fname = os.path.join(directory, fname)
     with open(fname, 'w') as f:
       json.dump(data, f)
+    return fname
 
   def GetRawData(self):
     """Getter for all raw_data."""
@@ -211,12 +219,16 @@ class StatsManager(object):
     Args:
       directory: directory to create the raw data folder in.
       dirname: folder in which raw data live.
+
+    Returns:
+      list of full path of each domain's raw data save location
     """
     if not os.path.exists(directory):
       os.makedirs(directory)
     dirname = os.path.join(directory, dirname)
     if not os.path.exists(dirname):
       os.makedirs(dirname)
+    fnames = []
     for domain, data in self._data.iteritems():
       if not domain.endswith(self._unit[domain]):
         domain = '%s_%s' % (domain, self._unit[domain])
@@ -224,3 +236,5 @@ class StatsManager(object):
       fname = os.path.join(dirname, fname)
       with open(fname, 'w') as f:
         f.write('\n'.join('%.2f' % sample for sample in data) + '\n')
+      fnames.append(fname)
+    return fnames
diff --git a/extra/usb_power/stats_manager_unittest.py b/extra/usb_power/stats_manager_unittest.py
index 3122fda455..2fb1080f4d 100644
--- a/extra/usb_power/stats_manager_unittest.py
+++ b/extra/usb_power/stats_manager_unittest.py
@@ -115,25 +115,31 @@ class TestStatsManager(unittest.TestCase):
     """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)
-    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())
+    expected_files = set(['A_mW.txt', 'B_mV.txt'])
+    fnames = self.data.SaveRawData(self.tempdir, dirname)
+    files_returned = set([os.path.basename(f) for f in fnames])
+    # Assert that only the expected files got returned.
+    self.assertEqual(expected_files, files_returned)
+    # Assert that only the returned files are in the outdir.
+    self.assertEqual(set(os.listdir(os.path.join(self.tempdir, dirname))),
+                     files_returned)
+    for fname in fnames:
+      with open(fname, 'r') as f:
+        if 'A_mW' in fname:
+          self.assertEqual('99999.50', f.readline().strip())
+          self.assertEqual('100000.50', f.readline().strip())
+        if 'B_mV' in fname:
+          self.assertEqual('1.50', f.readline().strip())
+          self.assertEqual('2.50', f.readline().strip())
+          self.assertEqual('3.50', f.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))
+    files = self.data.SaveRawData(self.tempdir, outdir)
+    files = [os.path.basename(f) for f in files]
     # Verify nothing gets appended to domain for filename if no unit exists.
     self.assertIn('B.txt', files)
 
@@ -166,8 +172,13 @@ class TestStatsManager(unittest.TestCase):
     """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)
+    expected_fname = os.path.join(self.tempdir, fname)
+    fname = self.data.SaveSummary(self.tempdir, fname)
+    # Assert the reported fname is the same as the expected fname
+    self.assertEqual(expected_fname, fname)
+    # Assert only the reported fname is output (in the tempdir)
+    self.assertEqual(set([os.path.basename(fname)]),
+                     set(os.listdir(self.tempdir)))
     with open(fname, 'r') as f:
       self.assertEqual(
           '@@   NAME  COUNT       MEAN  STDDEV        MAX       MIN\n',
@@ -183,8 +194,13 @@ class TestStatsManager(unittest.TestCase):
     """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)
+    expected_fname = os.path.join(self.tempdir, fname)
+    fname = self.data.SaveSummaryJSON(self.tempdir, fname)
+    # Assert the reported fname is the same as the expected fname
+    self.assertEqual(expected_fname, fname)
+    # Assert only the reported fname is output (in the tempdir)
+    self.assertEqual(set([os.path.basename(fname)]),
+                     set(os.listdir(self.tempdir)))
     with open(fname, 'r') as f:
       summary = json.load(f)
       self.assertAlmostEqual(100000.0, summary['A']['mean'])
@@ -197,8 +213,7 @@ class TestStatsManager(unittest.TestCase):
     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)
+    fname = self.data.SaveSummaryJSON(self.tempdir, fname)
     with open(fname, 'r') as f:
       summary = json.load(f)
       self.assertEqual('blue', summary['A']['unit'])