#!/usr/bin/env python3
# Copyright (c) 2023 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import io
import os
import os.path
import sys
import tempfile
import unittest
import unittest.mock

ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT_DIR)

import reclient_metrics


class ReclientMetricsTest(unittest.TestCase):
  def test_is_googler(self):
    with unittest.mock.patch('subprocess.run') as run_mock:
      run_mock.return_value.returncode = 0
      run_mock.return_value.stdout = 'Logged in as abc@google.com.'
      self.assertTrue(reclient_metrics.is_googler())

    with unittest.mock.patch('subprocess.run') as run_mock:
      run_mock.return_value.returncode = 1
      self.assertFalse(reclient_metrics.is_googler())

    with unittest.mock.patch('subprocess.run') as run_mock:
      run_mock.return_value.returncode = 0
      run_mock.return_value.stdout = ''
      self.assertFalse(reclient_metrics.is_googler())

    with unittest.mock.patch('subprocess.run') as run_mock:
      run_mock.return_value.returncode = 0
      run_mock.return_value.stdout = 'Logged in as foo@example.com.'
      self.assertFalse(reclient_metrics.is_googler())

    with unittest.mock.patch('subprocess.run') as run_mock:
      self.assertTrue(reclient_metrics.is_googler({
          'is-googler': True,
      }))
      self.assertFalse(reclient_metrics.is_googler({
          'is-googler': False,
      }))
      run_mock.assert_not_called()

  def test_load_and_save_config(self):
    with tempfile.TemporaryDirectory() as tmpdir:
      reclient_metrics.CONFIG = os.path.join(tmpdir, 'reclient_metrics.cfg')
      with unittest.mock.patch('subprocess.run') as run_mock:
        run_mock.return_value.returncode = 0
        run_mock.return_value.stdout = 'Logged in as abc@google.com.'
        cfg1 = reclient_metrics.load_config()
        self.assertDictEqual(
            cfg1, {
                'is-googler': True,
                'countdown': 10,
                'version': reclient_metrics.VERSION,
            })
        reclient_metrics.save_config(cfg1)
        cfg2 = reclient_metrics.load_config()
        self.assertDictEqual(
            cfg2, {
                'is-googler': True,
                'countdown': 9,
                'version': reclient_metrics.VERSION,
            })
        run_mock.assert_called_once()

  def test_check_status(self):
    with tempfile.TemporaryDirectory() as tmpdir:
      reclient_metrics.CONFIG = os.path.join(tmpdir, 'reclient_metrics.cfg')
      with unittest.mock.patch('subprocess.run') as run_mock:
        run_mock.return_value.returncode = 0
        run_mock.return_value.stdout = 'Logged in as abc@google.com.'
        with unittest.mock.patch('sys.stdout',
                                 new=io.StringIO()) as stdout_mock:
          self.assertFalse(reclient_metrics.check_status("outdir"))
          self.assertIn("Your reclient metrics will", stdout_mock.getvalue())
          self.assertIn(
              os.path.join("outdir", ".reproxy_tmp", "logs", "rbe_metrics.txt"),
              stdout_mock.getvalue())
        run_mock.assert_called_once()

    with tempfile.TemporaryDirectory() as tmpdir:
      reclient_metrics.CONFIG = os.path.join(tmpdir, 'reclient_metrics.cfg')
      with unittest.mock.patch('subprocess.run') as run_mock:
        run_mock.return_value.returncode = 0
        run_mock.return_value.stdout = 'Logged in as abc@google.com.'
        for i in range(10):
          with unittest.mock.patch('sys.stdout',
                                   new=io.StringIO()) as stdout_mock:
            self.assertFalse(reclient_metrics.check_status("outdir"))
            self.assertIn("Your reclient metrics will", stdout_mock.getvalue())
            self.assertIn(
                os.path.join("outdir", ".reproxy_tmp", "logs",
                             "rbe_metrics.txt"), stdout_mock.getvalue())
            self.assertIn("you run autoninja another %d time(s)" % (10 - i),
                          stdout_mock.getvalue())
        with unittest.mock.patch('sys.stdout',
                                 new=io.StringIO()) as stdout_mock:
          self.assertTrue(reclient_metrics.check_status("outdir"))
          self.assertNotIn("Your reclient metrics will", stdout_mock.getvalue())
          self.assertNotIn(
              os.path.join("outdir", ".reproxy_tmp", "logs", "rbe_metrics.txt"),
              stdout_mock.getvalue())
        run_mock.assert_called_once()

    with tempfile.TemporaryDirectory() as tmpdir:
      reclient_metrics.CONFIG = os.path.join(tmpdir, 'reclient_metrics.cfg')
      with unittest.mock.patch('subprocess.run') as run_mock:
        run_mock.return_value.returncode = 0
        run_mock.return_value.stdout = 'Logged in as abc@example.com.'
        with unittest.mock.patch('sys.stdout',
                                 new=io.StringIO()) as stdout_mock:
          self.assertFalse(reclient_metrics.check_status("outdir"))
          self.assertNotIn("Your reclient metrics will", stdout_mock.getvalue())
          self.assertNotIn(
              os.path.join("outdir", ".reproxy_tmp", "logs", "rbe_metrics.txt"),
              stdout_mock.getvalue())
        run_mock.assert_called_once()

    with tempfile.TemporaryDirectory() as tmpdir:
      reclient_metrics.CONFIG = os.path.join(tmpdir, 'reclient_metrics.cfg')
      with unittest.mock.patch('subprocess.run') as run_mock:
        run_mock.return_value.returncode = 1
        run_mock.return_value.stdout = ''
        with unittest.mock.patch('sys.stdout',
                                 new=io.StringIO()) as stdout_mock:
          self.assertFalse(reclient_metrics.check_status("outdir"))
          self.assertNotIn("Your reclient metrics will", stdout_mock.getvalue())
          self.assertNotIn(
              os.path.join("outdir", ".reproxy_tmp", "logs", "rbe_metrics.txt"),
              stdout_mock.getvalue())
        run_mock.assert_called_once()

    with tempfile.TemporaryDirectory() as tmpdir:
      reclient_metrics.CONFIG = os.path.join(tmpdir, 'reclient_metrics.cfg')
      with unittest.mock.patch('subprocess.run') as run_mock:
        run_mock.return_value.returncode = 0
        run_mock.return_value.stdout = 'Logged in as abc@google.com.'
        reclient_metrics.main(["reclient_metrics.py", "opt-in"])
        with unittest.mock.patch('sys.stdout',
                                 new=io.StringIO()) as stdout_mock:
          self.assertTrue(reclient_metrics.check_status("outdir"))
          self.assertNotIn("Your reclient metrics will", stdout_mock.getvalue())
          self.assertNotIn(
              os.path.join("outdir", ".reproxy_tmp", "logs", "rbe_metrics.txt"),
              stdout_mock.getvalue())
        run_mock.assert_called_once()

    with tempfile.TemporaryDirectory() as tmpdir:
      reclient_metrics.CONFIG = os.path.join(tmpdir, 'reclient_metrics.cfg')
      with unittest.mock.patch('subprocess.run') as run_mock:
        run_mock.return_value.returncode = 0
        run_mock.return_value.stdout = 'Logged in as abc@google.com.'
        for i in range(3):
          with unittest.mock.patch('sys.stdout',
                                   new=io.StringIO()) as stdout_mock:
            self.assertFalse(reclient_metrics.check_status("outdir"))
            self.assertIn("Your reclient metrics will", stdout_mock.getvalue())
            self.assertIn(
                os.path.join("outdir", ".reproxy_tmp", "logs",
                             "rbe_metrics.txt"), stdout_mock.getvalue())
            self.assertIn("you run autoninja another %d time(s)" % (10 - i),
                          stdout_mock.getvalue())
        reclient_metrics.main(["reclient_metrics.py", "opt-in"])
        with unittest.mock.patch('sys.stdout',
                                 new=io.StringIO()) as stdout_mock:
          self.assertTrue(reclient_metrics.check_status("outdir"))
          self.assertNotIn("Your reclient metrics will", stdout_mock.getvalue())
          self.assertNotIn(
              os.path.join("outdir", ".reproxy_tmp", "logs", "rbe_metrics.txt"),
              stdout_mock.getvalue())
        run_mock.assert_called_once()

    with tempfile.TemporaryDirectory() as tmpdir:
      reclient_metrics.CONFIG = os.path.join(tmpdir, 'reclient_metrics.cfg')
      with unittest.mock.patch('subprocess.run') as run_mock:
        run_mock.return_value.returncode = 0
        run_mock.return_value.stdout = 'Logged in as abc@example.com.'
        with unittest.mock.patch('sys.stdout',
                                 new=io.StringIO()) as stdout_mock:
          self.assertFalse(reclient_metrics.check_status("outdir"))
          self.assertNotIn("Your reclient metrics will", stdout_mock.getvalue())
          self.assertNotIn(
              os.path.join("outdir", ".reproxy_tmp", "logs", "rbe_metrics.txt"),
              stdout_mock.getvalue())
        reclient_metrics.main(["reclient_metrics.py", "opt-in"])
        with unittest.mock.patch('sys.stdout',
                                 new=io.StringIO()) as stdout_mock:
          self.assertFalse(reclient_metrics.check_status("outdir"))
          self.assertNotIn("Your reclient metrics will", stdout_mock.getvalue())
          self.assertNotIn(
              os.path.join("outdir", ".reproxy_tmp", "logs", "rbe_metrics.txt"),
              stdout_mock.getvalue())
        run_mock.assert_called_once()

    with tempfile.TemporaryDirectory() as tmpdir:
      reclient_metrics.CONFIG = os.path.join(tmpdir, 'reclient_metrics.cfg')
      with unittest.mock.patch('subprocess.run') as run_mock:
        run_mock.return_value.returncode = 0
        run_mock.return_value.stdout = 'Logged in as abc@google.com.'
        for i in range(3):
          with unittest.mock.patch('sys.stdout',
                                   new=io.StringIO()) as stdout_mock:
            self.assertFalse(reclient_metrics.check_status("outdir"))
            self.assertIn("Your reclient metrics will", stdout_mock.getvalue())
            self.assertIn(
                os.path.join("outdir", ".reproxy_tmp", "logs",
                             "rbe_metrics.txt"), stdout_mock.getvalue())
            self.assertIn("you run autoninja another %d time(s)" % (10 - i),
                          stdout_mock.getvalue())
        reclient_metrics.main(["reclient_metrics.py", "opt-out"])
        with unittest.mock.patch('sys.stdout',
                                 new=io.StringIO()) as stdout_mock:
          self.assertFalse(reclient_metrics.check_status("outdir"))
          self.assertNotIn("Your reclient metrics will", stdout_mock.getvalue())
          self.assertNotIn(
              os.path.join("outdir", ".reproxy_tmp", "logs", "rbe_metrics.txt"),
              stdout_mock.getvalue())
        run_mock.assert_called_once()


if __name__ == '__main__':
  unittest.main()