diff --git a/metrics.py b/metrics.py index 714ebadac..2132b027d 100644 --- a/metrics.py +++ b/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 diff --git a/tests/metrics_test.py b/tests/metrics_test.py index a4955162e..fa4cbc604 100644 --- a/tests/metrics_test.py +++ b/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()