# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for lab_cli."""
import atexit
import copy
import datetime
import json
import os
import shutil
import tempfile
from unittest import mock

from absl.testing import absltest
from absl.testing import parameterized
from tradefed_cluster.configs import lab_config


from multitest_transport.cli import lab_cli
from multitest_transport.cli import unittest_util


_PROJECT = 'aproject'
_SERVICE_ACCOUNT_EMAIL = 'as@gsa.com'
_KEY_CREATE_TIME = '2021-04-15T23:13:34Z'
_SERVICE_ACCOUNT_KEY_ID = 'akey_id'
_SERVICE_ACCOUNT_KEY = {
    'type': 'service_account',
    'project_id': _PROJECT,
    'private_key_id': _SERVICE_ACCOUNT_KEY_ID,
    'private_key': 'aprivate_key',
    'client_email': _SERVICE_ACCOUNT_EMAIL,
    'client_id': '1234',
    'auth_uri': 'https://accounts.google.com/o/oauth2/auth',
    'token_uri': 'https://oauth2.googleapis.com/token',
    'auth_provider_x509_cert_url': 'https://www.googleapis.com/oauth2/v1/certs',
    'client_x509_cert_url': 'url'
}
_SERVICE_ACCOUNT_KEY_INFO = {
    'name': (f'projects/{_PROJECT}/serviceAccounts/'
             f'{_SERVICE_ACCOUNT_EMAIL}/keys/{_SERVICE_ACCOUNT_KEY_ID}'),
    'validAfterTime': _KEY_CREATE_TIME,
    'validBeforeTime': '2021-07-14T23:13:34Z',
    'keyAlgorithm': 'KEY_ALG_RSA_2048',
    'keyOrigin': 'GOOGLE_PROVIDED',
    'keyType': 'USER_MANAGED'
}


class LabCliTest(parameterized.TestCase):

  def setUp(self):
    super(LabCliTest, self).setUp()
    self.arg_parser = lab_cli.CreateParser()
    self.tmp_root = tempfile.mkdtemp()

  def tearDown(self):
    shutil.rmtree(self.tmp_root)
    super(LabCliTest, self).tearDown()

  def testCreateParser(self):
    args = self.arg_parser.parse_args(['version'])
    self.assertEqual('version', args.action)

  def testCreateParser_withLoggingArgs(self):
    args = self.arg_parser.parse_args([
        '-v', '--log_file', 'file.log', 'version'])
    self.assertEqual('version', args.action)
    self.assertTrue(args.verbose)
    self.assertEqual('file.log', args.log_file)

  @mock.patch.object(tempfile, 'mkdtemp')
  def testSetupMTTBinary(self, mock_create_temp):
    tmp_folder = os.path.join(self.tmp_root, 'mtt_lab_extracted')
    os.mkdir(tmp_folder)
    mock_create_temp.return_value = tmp_folder
    mtt_lab_path = os.path.join(self.tmp_root, 'mtt_lab')
    unittest_util.CreateZipFile(
        mtt_lab_path,
        [unittest_util.File(
            filename='mtt_binary', content='this is a mtt binary')])

    args = mock.MagicMock(cli_path=mtt_lab_path)
    host = mock.MagicMock(
        context=mock.MagicMock(user='testuser'))
    lab_cli._SetupMTTBinary(args, host)

    host.context.assert_has_calls([
        mock.call.CopyFile(
            os.path.join(tmp_folder, 'mtt_binary'),
            '/tmp/testuser/mtt'),
        mock.call.Run(['chmod', '+x', '/tmp/testuser/mtt'])])
    self.assertEqual('Setting up MTT Binary', host.execution_state)

  @mock.patch.object(tempfile, 'NamedTemporaryFile')
  def testSetupHostConfig(self, mock_make_tmpfile):
    mock_tmpfile = mock.MagicMock()
    mock_tmpfile.name = '/local/config.yaml'
    mock_make_tmpfile.return_value = mock_tmpfile
    host = mock.MagicMock(context=mock.MagicMock(user='testuser'))
    lab_cli._SetupHostConfig(host)

    host.config.Save.assert_called_once_with('/local/config.yaml')
    host.context.CopyFile.assert_called_once_with(
        '/local/config.yaml', '/tmp/testuser/mtt_host_config.yaml')
    mock_tmpfile.close.assert_called_once_with()
    self.assertEqual('Setting up host config', host.execution_state)

  @parameterized.named_parameters(
      ('$ mtt_lab start --ask_sudo_password --sudo_user user',
       True, 'user', True),
      ('$ mtt_lab start --sudo_user user', False, 'user', True),
      ('$ mtt_lab start --ask_sudo_password', True, None, True),
      ('$ mtt_lab start', False, None, False))
  @mock.patch.object(lab_cli, '_SetupMTTBinary')
  @mock.patch.object(tempfile, 'NamedTemporaryFile')
  @mock.patch.object(lab_config.HostConfig, 'Save')
  def testStart(
      self, ask_sudo_password, sudo_user, use_sudo, mock_save,
      mock_make_tmpfile, mock_setup):
    mock_tmpfile = mock.MagicMock()
    mock_tmpfile.name = '/local/config.yaml'
    mock_make_tmpfile.return_value = mock_tmpfile
    args = mock.MagicMock(
        service_account_json_key_path='/path/to/keyfile.json',
        ask_sudo_password=ask_sudo_password,
        sudo_user=sudo_user,
        verbose=False,
        very_verbose=False)
    host = mock.MagicMock(
        context=mock.MagicMock(user='testuser'),
        config=lab_config.CreateHostConfig(
            service_account_json_key_path='/path/to/keyfile.json'))
    lab_cli.Start(args, host)

    host.context.assert_has_calls([
        mock.call.CopyFile('/path/to/keyfile.json',
                           '/tmp/testuser/keyfile/key.json'),
        mock.call.CopyFile('/local/config.yaml',
                           '/tmp/testuser/mtt_host_config.yaml'),
        mock.call.Run(['/tmp/testuser/mtt', '--no_check_update', 'start',
                       '/tmp/testuser/mtt_host_config.yaml'],
                      sudo=use_sudo)])
    self.assertEqual('/tmp/testuser/keyfile/key.json',
                     host.config.service_account_json_key_path)
    mock_save.assert_called_once_with('/local/config.yaml')
    mock_setup.assert_called_once_with(args, host)
    mock_tmpfile.close.assert_called_once_with()
    self.assertEqual('Running start', host.execution_state)

  @mock.patch.object(lab_cli, '_SetupMTTBinary')
  @mock.patch.object(tempfile, 'NamedTemporaryFile')
  @mock.patch.object(lab_config.HostConfig, 'Save')
  def testStart_noServiceAccountKey(
      self, mock_save, mock_make_tmpfile, mock_setup):
    mock_tmpfile = mock.MagicMock()
    mock_tmpfile.name = '/local/config.yaml'
    mock_make_tmpfile.return_value = mock_tmpfile
    args = mock.MagicMock(
        service_account_json_key_path=None,
        sudo_user=None,
        ask_sudo_password=False,
        verbose=False,
        very_verbose=False)
    host = mock.MagicMock(
        context=mock.MagicMock(user='testuser'),
        config=lab_config.CreateHostConfig(
            service_account_json_key_path=None))
    lab_cli.Start(args, host)

    host.context.assert_has_calls([
        mock.call.CopyFile(
            '/local/config.yaml', '/tmp/testuser/mtt_host_config.yaml'),
        mock.call.Run(['/tmp/testuser/mtt', '--no_check_update', 'start',
                       '/tmp/testuser/mtt_host_config.yaml'], sudo=False)])
    mock_save.assert_called_once_with('/local/config.yaml')
    mock_setup.assert_called_once_with(args, host)
    mock_tmpfile.close.assert_called_once_with()

  @parameterized.named_parameters(
      ('$ mtt_lab update --ask_sudo_password --sudo_user user',
       True, 'user', True),
      ('$ mtt_lab update --sudo_user user', False, 'user', True),
      ('$ mtt_lab update --ask_sudo_password', True, None, True),
      ('$ mtt_lab update', False, None, False))
  @mock.patch.object(lab_cli, '_SetupMTTBinary')
  @mock.patch.object(tempfile, 'NamedTemporaryFile')
  @mock.patch.object(lab_config.HostConfig, 'Save')
  def testUpdate(
      self, ask_sudo_password, sudo_user, use_sudo, mock_save,
      mock_make_tmpfile, mock_setup):
    mock_tmpfile = mock.MagicMock()
    mock_tmpfile.name = '/local/config.yaml'
    mock_make_tmpfile.return_value = mock_tmpfile
    args = mock.MagicMock(
        service_account_json_key_path='/path/to/keyfile.json',
        ask_sudo_password=ask_sudo_password,
        sudo_user=sudo_user,
        verbose=False,
        very_verbose=False)
    host = mock.MagicMock(
        context=mock.MagicMock(user='testuser'),
        config=lab_config.CreateHostConfig(
            service_account_json_key_path='/path/to/keyfile.json'))
    lab_cli.Update(args, host)

    host.context.assert_has_calls([
        mock.call.CopyFile('/path/to/keyfile.json',
                           '/tmp/testuser/keyfile/key.json'),
        mock.call.CopyFile('/local/config.yaml',
                           '/tmp/testuser/mtt_host_config.yaml'),
        mock.call.Run(['/tmp/testuser/mtt', '--no_check_update', 'update',
                       '/tmp/testuser/mtt_host_config.yaml'],
                      sudo=use_sudo)])
    self.assertEqual('/tmp/testuser/keyfile/key.json',
                     host.config.service_account_json_key_path)
    mock_save.assert_called_once_with('/local/config.yaml')
    mock_setup.assert_called_once_with(args, host)
    mock_tmpfile.close.assert_called_once_with()
    self.assertEqual('Running update', host.execution_state)

  @parameterized.named_parameters(
      ('$ mtt_lab restart --ask_sudo_password --sudo_user user',
       True, 'user', True),
      ('$ mtt_lab restart --sudo_user user', False, 'user', True),
      ('$ mtt_lab restart --ask_sudo_password', True, None, True),
      ('$ mtt_lab restart', False, None, False))
  @mock.patch.object(lab_cli, '_SetupMTTBinary')
  @mock.patch.object(tempfile, 'NamedTemporaryFile')
  @mock.patch.object(lab_config.HostConfig, 'Save')
  def testRestart(
      self, ask_sudo_password, sudo_user, use_sudo, mock_save,
      mock_make_tmpfile, mock_setup):
    mock_tmpfile = mock.MagicMock()
    mock_tmpfile.name = '/local/config.yaml'
    mock_make_tmpfile.return_value = mock_tmpfile
    args = mock.MagicMock(
        service_account_json_key_path='/path/to/keyfile.json',
        ask_sudo_password=ask_sudo_password,
        sudo_user=sudo_user,
        verbose=False,
        very_verbose=False)
    host = mock.MagicMock(
        context=mock.MagicMock(user='testuser'),
        config=lab_config.CreateHostConfig(
            service_account_json_key_path='/path/to/keyfile.json'))
    lab_cli.Restart(args, host)

    host.context.assert_has_calls([
        mock.call.CopyFile('/path/to/keyfile.json',
                           '/tmp/testuser/keyfile/key.json'),
        mock.call.CopyFile('/local/config.yaml',
                           '/tmp/testuser/mtt_host_config.yaml'),
        mock.call.Run(['/tmp/testuser/mtt', '--no_check_update', 'restart',
                       '/tmp/testuser/mtt_host_config.yaml'],
                      sudo=use_sudo)])
    self.assertEqual('/tmp/testuser/keyfile/key.json',
                     host.config.service_account_json_key_path)
    mock_save.assert_called_once_with('/local/config.yaml')
    mock_setup.assert_called_once_with(args, host)
    mock_tmpfile.close.assert_called_once_with()
    self.assertEqual('Running restart', host.execution_state)

  @parameterized.named_parameters(
      ('$ mtt_lab stop --ask_sudo_password --sudo_user user',
       True, 'user', True),
      ('$ mtt_lab stop --sudo_user user', False, 'user', True),
      ('$ mtt_lab stop --ask_sudo_password', True, None, True),
      ('$ mtt_lab stop', False, None, False))
  @mock.patch.object(lab_cli, '_SetupMTTBinary')
  def testStop(self, ask_sudo_password, sudo_user, use_sudo, mock_setup):
    args = mock.MagicMock(
        ask_sudo_password=ask_sudo_password,
        sudo_user=sudo_user,
        verbose=False,
        very_verbose=False)
    host = mock.MagicMock(context=mock.MagicMock(user='testuser'))
    lab_cli.Stop(args, host)

    host.context.Run.assert_called_once_with(
        ['/tmp/testuser/mtt', '--no_check_update', 'stop'],
        sudo=use_sudo)
    mock_setup.assert_called_once_with(args, host)
    self.assertEqual('Running stop', host.execution_state)

  @parameterized.named_parameters(
      ('run_cmd_with_sudo_pwd_and_user', True, 'user', True),
      ('run_cmd_with_sudo_user', False, 'user', True),
      ('run_cmd_with_sudo_pwd', True, 'user', True),
      ('run_cmd_without_sudo', False, None, False))
  def testRunCmd(self, ask_sudo_password, sudo_user, use_sudo):
    args = mock.MagicMock(
        ask_sudo_password=ask_sudo_password,
        sudo_user=sudo_user,
        cmd='run a command line')
    host = mock.MagicMock()
    lab_cli.RunCmd(args, host)
    host.context.Run.assert_called_once_with(
        ['run', 'a', 'command', 'line'], sudo=use_sudo)
    self.assertEqual('Running cmd', host.execution_state)

  def testCreateLabCommandArgParser(self):
    parser = lab_cli._CreateLabCommandArgParser()
    args = parser.parse_args(['lab_config.yaml'])
    self.assertEqual('lab_config.yaml', args.lab_config_path)
    self.assertEmpty(args.hosts_or_clusters)

  def testCreateLabCommandArgParser_withHostOrCluster(self):
    parser = lab_cli._CreateLabCommandArgParser()
    args = parser.parse_args(['lab_config.yaml', 'host1', 'cluster1'])
    self.assertEqual('lab_config.yaml', args.lab_config_path)
    self.assertEqual(['host1', 'cluster1'], args.hosts_or_clusters)

  @parameterized.named_parameters(
      ('no_verbose', [], ['--no_check_update']),
      ('verbose', ['-v'], ['-v', '--no_check_update']),
      ('very_verbose', ['-vv'], ['-vv', '--no_check_update']),
      ('both', ['-v', '-vv'], ['-vv', '--no_check_update']))
  def testBuildBaseMTTCmd(self, verbose_flag, res_flag):
    parser = lab_cli.cli_util.CreateLoggingArgParser()
    args = parser.parse_args(verbose_flag)
    host = mock.MagicMock(context=mock.MagicMock(user='testuser'))
    cmd = lab_cli._BuildBaseMTTCmd(args, host)
    self.assertEqual(
        ['/tmp/testuser/mtt'] + res_flag,
        cmd)

  def testSetupServiceAccountKey(self):
    args = mock.MagicMock(service_account_json_key_path=None)
    host = mock.MagicMock(
        context=mock.MagicMock(user='testuser'),
        config=lab_config.CreateHostConfig(
            service_account_json_key_path='/path/to/keyfile.json'))
    lab_cli._SetupServiceAccountKey(args, host)

    host.context.CopyFile.assert_called_once_with(
        '/path/to/keyfile.json', '/tmp/testuser/keyfile/key.json')
    self.assertEqual('/tmp/testuser/keyfile/key.json',
                     host.config.service_account_json_key_path)
    self.assertEqual('Setting up service account key', host.execution_state)

  @mock.patch.object(tempfile, 'NamedTemporaryFile')
  @mock.patch.object(lab_cli.google_auth_util, 'CreateKey')
  @mock.patch.object(lab_cli.google_auth_util, 'UpdateSecret')
  @mock.patch.object(lab_cli.google_auth_util, 'GetServiceAccountKeyInfo')
  @mock.patch.object(lab_cli.google_auth_util, 'GetSecret')
  @mock.patch.object(lab_cli.google_auth_util, 'CanCreateKey')
  @mock.patch.object(lab_cli.google_auth_util, 'CanUpdateSecret')
  @mock.patch.object(lab_cli.google_auth_util, 'DisableSecretVersions')
  @mock.patch.object(atexit, 'register')
  def testGetServiceAccountKeyFilePath(
      self, mock_register, mock_disable_versions, mock_can_update_secret,
      mock_can_create_key, mock_get_secret, mock_get_key_info,
      mock_update_secret, mock_create_key, mock_make_tmpfile):
    mock_tmpfile = mock.MagicMock()
    mock_tmpfile.name = '/local/sa_key.json'
    mock_make_tmpfile.return_value = mock_tmpfile
    mock_get_secret.return_value = json.dumps(_SERVICE_ACCOUNT_KEY).encode()
    key_info = copy.deepcopy(_SERVICE_ACCOUNT_KEY_INFO)
    key_info['validAfterTime'] = (
        datetime.datetime.now(tz=datetime.timezone.utc)
        - datetime.timedelta(days=3)).isoformat()
    mock_get_key_info.return_value = key_info

    lab_cli._GetServiceAccountKeyFilePath('secret_project', 'sa_key')

    mock_get_key_info.assert_called_once_with(
        _SERVICE_ACCOUNT_EMAIL, _SERVICE_ACCOUNT_KEY_ID)
    self.assertTrue(mock_can_create_key.called)
    self.assertTrue(mock_can_update_secret.called)
    self.assertFalse(mock_create_key.called)
    self.assertFalse(mock_update_secret.called)
    self.assertFalse(mock_disable_versions.called)
    mock_tmpfile.write.assert_called_once_with(
        json.dumps(_SERVICE_ACCOUNT_KEY).encode())
    mock_tmpfile.flush.assert_called_once_with()
    mock_get_secret.assert_called_once_with('secret_project', 'sa_key')
    mock_register.assert_called_once_with(mock_tmpfile.close)

  @mock.patch.object(tempfile, 'NamedTemporaryFile')
  @mock.patch.object(lab_cli.google_auth_util, 'CreateKey')
  @mock.patch.object(lab_cli.google_auth_util, 'UpdateSecret')
  @mock.patch.object(lab_cli.google_auth_util, 'GetServiceAccountKeyInfo')
  @mock.patch.object(lab_cli.google_auth_util, 'GetSecret')
  @mock.patch.object(lab_cli.google_auth_util, 'CanCreateKey')
  @mock.patch.object(lab_cli.google_auth_util, 'CanUpdateSecret')
  @mock.patch.object(lab_cli.google_auth_util, 'DisableSecretVersions')
  @mock.patch.object(atexit, 'register')
  def testGetServiceAccountKeyFilePath_renew(
      self, mock_register, mock_disable_versions, mock_can_update_secret,
      mock_can_create_key, mock_get_secret, mock_get_key_info,
      mock_update_secret, mock_create_key, mock_make_tmpfile):
    mock_tmpfile = mock.MagicMock()
    mock_tmpfile.name = '/local/sa_key.json'
    mock_make_tmpfile.return_value = mock_tmpfile
    mock_get_secret.return_value = json.dumps(_SERVICE_ACCOUNT_KEY).encode()
    key_info = copy.deepcopy(_SERVICE_ACCOUNT_KEY_INFO)
    key_info['validAfterTime'] = (
        datetime.datetime.now(tz=datetime.timezone.utc)
        - datetime.timedelta(days=20)).isoformat()
    mock_get_key_info.return_value = key_info
    mock_can_create_key.return_value = True
    mock_can_update_secret.return_value = True
    new_key = copy.deepcopy(_SERVICE_ACCOUNT_KEY)
    new_key['private_key_id'] = 'new_key_id'
    mock_create_key.return_value = new_key
    mock_update_secret.return_value = 'new_version'

    lab_cli._GetServiceAccountKeyFilePath('secret_project', 'sa_key')

    mock_get_key_info.assert_called_once_with(
        _SERVICE_ACCOUNT_EMAIL, _SERVICE_ACCOUNT_KEY_ID)
    mock_can_update_secret.assert_called_once_with('secret_project', 'sa_key')
    mock_can_create_key.assert_called_once_with(_SERVICE_ACCOUNT_EMAIL)
    mock_create_key.assert_called_once_with(_SERVICE_ACCOUNT_EMAIL)
    mock_update_secret.assert_called_once_with(
        'secret_project', 'sa_key', json.dumps(new_key).encode())
    mock_disable_versions.assert_called_once_with(
        'secret_project', 'sa_key', exclude_versions=['new_version'],
        days_before=28)
    mock_tmpfile.write.assert_called_once_with(json.dumps(new_key).encode())
    mock_tmpfile.flush.assert_called_once_with()
    mock_get_secret.assert_called_once_with('secret_project', 'sa_key')
    mock_register.assert_called_once_with(mock_tmpfile.close)

  @mock.patch.object(lab_cli.google_auth_util, 'GetServiceAccountKeyInfo')
  def testShouldRenewServiceAccountKey_noNeedToRenew(self, mock_get_key_info):
    key_info = copy.deepcopy(_SERVICE_ACCOUNT_KEY_INFO)
    key_info['validAfterTime'] = (
        datetime.datetime.now(tz=datetime.timezone.utc)
        - datetime.timedelta(days=3)).isoformat()
    mock_get_key_info.return_value = key_info

    shoud_renew = lab_cli._ShouldRenewServiceAccountKey(_SERVICE_ACCOUNT_KEY)

    self.assertFalse(shoud_renew)
    mock_get_key_info.assert_called_once_with(
        _SERVICE_ACCOUNT_EMAIL, _SERVICE_ACCOUNT_KEY_ID)

  @mock.patch.object(lab_cli.google_auth_util, 'GetServiceAccountKeyInfo')
  def testShouldRenewServiceAccountKey_shouldRenew(self, mock_get_key_info):
    key_info = copy.deepcopy(_SERVICE_ACCOUNT_KEY_INFO)
    key_info['validAfterTime'] = (
        datetime.datetime.now(tz=datetime.timezone.utc)
        - datetime.timedelta(days=20)).isoformat()
    mock_get_key_info.return_value = key_info

    shoud_renew = lab_cli._ShouldRenewServiceAccountKey(_SERVICE_ACCOUNT_KEY)

    self.assertTrue(shoud_renew)
    mock_get_key_info.assert_called_once_with(
        _SERVICE_ACCOUNT_EMAIL, _SERVICE_ACCOUNT_KEY_ID)


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