metrics: Fix countdown to metrics collection.

We were printing the notice when the decorator was called, so that if
several functions are decorated, the countdown will increase more than once
during each execution.

This change fixes that.

Bug: 832386
Change-Id: I5d6dd2d793137e1e2c83b6cc765b6245080b9d91
Reviewed-on: https://chromium-review.googlesource.com/1141073
Reviewed-by: Aaron Gable <agable@chromium.org>
Commit-Queue: Edward Lesmes <ehmaldonado@chromium.org>

metrics.py

@@ -133,6 +133,15 @@ class MetricsCollector(object):
     p.stdin.write(json.dumps(self._reported_metrics))
 
   def _collect_metrics(self, func, command_name, *args, **kwargs):
+    # If the user hasn't opted in or out, and the countdown is not yet 0, just
+    # display the notice.
+    if self.config.opted_in == None and self.config.countdown > 0:
+      metrics_utils.print_notice(self.config.countdown)
+      self.config.decrease_countdown()
+      func(*args, **kwargs)
+      return
+
+    self._collecting_metrics = True
     self.add('command', command_name)
     try:
       start = time.time()
@@ -183,18 +192,11 @@ class MetricsCollector(object):
       # need to do anything.
       if self.config.opted_in == False or not self.config.is_googler:
         return func
-      # If the user hasn't opted in or out, and the countdown is not yet 0, just
-      # display the notice.
-      if self.config.opted_in == None and self.config.countdown > 0:
-        metrics_utils.print_notice(self.config.countdown)
-        self.config.decrease_countdown()
-        return func
       # Otherwise, collect the metrics.
       # Needed to preserve the __name__ and __doc__ attributes of func.
       @functools.wraps(func)
       def _inner(*args, **kwargs):
         self._collect_metrics(func, command_name, *args, **kwargs)
-      self._collecting_metrics = True
       return _inner
     return _decorator
 

tests/metrics_test.py

@@ -28,6 +28,7 @@ class TimeMock(object):
 
 class MetricsCollectorTest(unittest.TestCase):
   def setUp(self):
+    self.config_file = os.path.join(ROOT_DIR, 'metrics.cfg')
     self.collector = metrics.MetricsCollector()
 
     # Keep track of the URL requests, file reads/writes and subprocess spawned.
@@ -295,6 +296,38 @@ class MetricsCollectorTest(unittest.TestCase):
     self.assertEqual(cm.exception.code, 123)
     self.assert_collects_metrics({'exit_code': 123})
 
+  def test_counts_down(self):
+    """Tests that the countdown works correctly."""
+    self.FileRead.side_effect = [
+        '{"is-googler": true, "countdown": 10, "opt-in": null}'
+    ]
+
+    # We define multiple functions to ensure it has no impact on countdown.
+    @self.collector.collect_metrics('fun')
+    def fun():
+      pass
+    @self.collector.collect_metrics('foon')
+    def _foon():
+      pass
+
+    # Assert that the countdown hasn't decrease yet.
+    self.assertFalse(self.FileWrite.called)
+    self.assertEqual(self.collector.config.countdown, 10)
+
+    fun()
+
+    # Assert that the countdown decreased by one, and the config file was
+    # updated.
+    self.assertEqual(self.collector.config.countdown, 9)
+    self.print_notice.assert_called_once_with(10)
+
+    self.assertEqual(len(self.FileWrite.mock_calls), 1)
+    config_file, config = self.FileWrite.mock_calls[0][1]
+
+    self.assertEqual(config_file, self.config_file)
+    self.assertEqual(json.loads(config),
+                     {'is-googler': True, 'countdown': 9, 'opt-in': None})
+
 
 if __name__ == '__main__':
   unittest.main()