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='******'),
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()
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='******'),
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)
def testUpdateHostConfig(self):
host_config = lab_config.CreateHostConfig(
lab_name='alab',
cluster_name='acluster',
hostname='ahost',
host_login_name='auser',
tf_global_config_path='apath',
docker_image='a_docker_image',
graceful_shutdown=True,
enable_stackdriver=True,
enable_autoupdate=True,
service_account_json_key_path='a_service_keyfile',
control_server_url='tfc',
max_local_virtual_devices=1)
new_host_config = host_config.SetDockerImage('b_docker_image')
new_host_config = new_host_config.SetServiceAccountJsonKeyPath(
'b_service_keyfile')
self.assertEqual('a_service_keyfile',
host_config.service_account_json_key_path)
self.assertEqual('b_service_keyfile',
new_host_config.service_account_json_key_path)
self.assertEqual('a_docker_image', host_config.docker_image)
self.assertEqual('b_docker_image', new_host_config.docker_image)
self.assertEqual('alab', new_host_config.lab_name)
self.assertEqual('acluster', new_host_config.cluster_name)
self.assertEqual(1, new_host_config.max_local_virtual_devices)
def testCreateHostConfig_empty(self):
host_config = lab_config.CreateHostConfig()
self.assertEqual('', host_config.lab_name)
self.assertEqual('', host_config.cluster_name)
self.assertEqual('', host_config.hostname)
self.assertEqual('', host_config.host_login_name)
self.assertEqual('', host_config.tf_global_config_path)
def testCreateHostConfig_noLabName(self):
host_config = lab_config.CreateHostConfig(cluster_name='acluster',
hostname='ahost',
host_login_name='auser')
self.assertEqual('', host_config.lab_name)
self.assertEqual('acluster', host_config.cluster_name)
self.assertEqual('ahost', host_config.hostname)
self.assertEqual('auser', host_config.host_login_name)
def testSetupServiceAccountKey(self):
args = mock.MagicMock(service_account_json_key_path=None)
host = mock.MagicMock(
context=mock.MagicMock(user='******'),
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)
def testConfigEquals_notEquals(self):
host_config_1 = lab_config.CreateHostConfig(
lab_name='alab',
cluster_name='acluster',
hostname='ahost',
host_login_name='auser',
tf_global_config_path='apath',
docker_server='a_docker_server',
docker_image='a_docker_image',
graceful_shutdown=True,
enable_stackdriver=True,
enable_autoupdate=True,
service_account_json_key_path='a_service_keyfile',
control_server_url='tfc')
host_config_2 = host_config_1.SetDockerImage('b_docker_image')
self.assertNotEqual(host_config_1, host_config_2)
def testCreateHostConfig_withTmpfsConfigs(self):
host_config = lab_config.CreateHostConfig(
'alab',
'acluster',
'ahost',
'auser',
'apath',
tmpfs_configs=[
lab_config.CreateTmpfsConfig('/atmpfs', 1000, '750')
])
self.assertEqual('alab', host_config.lab_name)
self.assertEqual('acluster', host_config.cluster_name)
self.assertEqual('ahost', host_config.hostname)
self.assertEqual('auser', host_config.host_login_name)
self.assertEqual('apath', host_config.tf_global_config_path)
self.assertEqual(1, len(host_config.tmpfs_configs))
self.assertEqual('/atmpfs', host_config.tmpfs_configs[0].path)
self.assertEqual(1000, host_config.tmpfs_configs[0].size)
self.assertEqual('750', host_config.tmpfs_configs[0].mode)
def CreateHost(args):
"""Create a Host object for local host."""
key_path = getattr(args, 'service_account_json_key_path', None)
if not args.lab_config_path:
# MTT standalone mode
logger.info('No lab config path set; using standalone mode config')
host_config = lab_config.CreateHostConfig(
cluster_name='default',
hostname=socket.getfqdn(),
docker_image=_DEFAULT_MTT_IMAGE,
shutdown_timeout_sec=60)
else:
host_config = _GetHostConfig(
lab_config_path=args.lab_config_path, key_path=key_path)
if not host_config.docker_image:
host_config = host_config.SetDockerImage(_DEFAULT_DOCKERIZED_TF_IMAGE)
host = Host(host_config, context=command_util.CommandContext())
# override host configs from input args.
if key_path:
host.config = host.config.SetServiceAccountJsonKeyPath(key_path)
return host
def testCreateHostConfig(self):
host_config = lab_config.CreateHostConfig(
lab_name='alab',
cluster_name='acluster',
hostname='ahost',
host_login_name='auser',
tf_global_config_path='apath',
docker_server='a_docker_server',
docker_image='a_docker_image',
graceful_shutdown=True,
shutdown_timeout_sec=240,
enable_stackdriver=True,
enable_autoupdate=True,
extra_docker_args=['--arg1', 'value1'],
control_server_url='tfc',
secret_project_id='secret_project',
service_account_key_secret_id='sa_key',
service_account='*****@*****.**',
operation_mode='ON_PREMISE')
self.assertEqual('alab', host_config.lab_name)
self.assertEqual('acluster', host_config.cluster_name)
self.assertEqual('ahost', host_config.hostname)
self.assertEqual('auser', host_config.host_login_name)
self.assertEqual('apath', host_config.tf_global_config_path)
self.assertEqual('a_docker_image', host_config.docker_image)
self.assertEqual('tfc', host_config.control_server_url)
self.assertEqual('a_docker_server', host_config.docker_server)
self.assertTrue(host_config.graceful_shutdown)
self.assertEqual(240, host_config.shutdown_timeout_sec)
self.assertTrue(host_config.enable_stackdriver)
self.assertTrue(host_config.enable_autoupdate)
self.assertEqual(['--arg1', 'value1'], host_config.extra_docker_args)
self.assertEqual('secret_project', host_config.secret_project_id)
self.assertEqual('sa_key', host_config.service_account_key_secret_id)
self.assertEqual('*****@*****.**', host_config.service_account)
self.assertEqual(lab_config_pb2.OperationMode.ON_PREMISE,
host_config.operation_mode)
