def testCreateCommandContext_withSudo(self):
ssh_config = ssh_util.SshConfig(hostname='host1',
user='******',
password='******')
sudo_ssh_config = ssh_util.SshConfig(hostname='host1',
user='******',
password='******')
context = command_util.CommandContext(host='host1',
user='******',
ssh_config=ssh_config,
sudo_ssh_config=sudo_ssh_config)
self.assertIsNotNone(context)
self.assertFalse(context.IsLocal())
self.assertFalse(self.mock_ssh_context_creator.called)
self.mock_fabric_pkg.Connection.assert_has_calls([
mock.call(host='host1',
user='******',
config=mock.ANY,
connect_kwargs={'password': '******'}),
mock.call().open(),
mock.call(host='host1',
user='******',
config=mock.ANY,
connect_kwargs={'password': '******'}),
mock.call().open(),
mock.call().sudo('echo "Testing sudo access..."',
hide=True,
password='******')
])
def BuildLabConfigPool(lab_config_path, key_path=None):
"""Build lab config pool based on configured config path."""
if not lab_config_path:
logger.debug('No lab config path set.')
lab_config_pool = lab_config.LabConfigPool()
elif lab_config_path.startswith('gs:'):
logger.debug('Use lab config path %s.', lab_config_path)
if key_path:
cred = google_auth_util.CreateCredentialFromServiceAccount(
key_path,
scopes=[google_auth_util.GCS_READ_SCOPE])
else:
cred = google_auth_util.GetGCloudCredential(
command_util.CommandContext())
gcs_client = gcs_file_util.CreateGCSClient(_MTT_PROJECT, cred)
lab_config_pool = lab_config.LabConfigPool(
gcs_file_util.GCSFileEnumerator(
gcs_client, lab_config_path, lab_config.IsYaml))
else:
logger.debug('Use lab config path %s.', lab_config_path)
lab_config_pool = lab_config.LabConfigPool(
lab_config.LocalFileEnumerator(
lab_config_path, lab_config.IsYaml))
lab_config_pool.LoadConfigs()
return lab_config_pool
def Create(self, args):
"""Creates a cluster.
This actually creates a Docker swarm and deploy a MTT service on it.
Args:
args: an argparse.ArgumentParser object.
Raises:
ValueError: if mtt_control_server_url or host is not set.
"""
if not config.config.mtt_control_server_url:
raise ValueError('mtt_control_server_url must be set.')
if not args.host:
raise ValueError('--host option must be set')
context = command_util.CommandContext(host=args.host,
user=args.ssh_user)
docker_context = command_util.DockerContext(context,
try_use_gcloud=False)
cluster_config = self._registry.GetConfig(args.name)
docker_context.Run(['swarm', 'init'])
# TODO: get token ID and store it.
docker_context.Run([
'service', 'create', '--name', 'mtt', '--env',
'MTT_CONTROL_SERVER_URL=%s' % config.config.mtt_control_server_url,
'--mode', 'global', 'gcr.io/android-mtt/mtt'
])
cluster_config.manager_host = args.host
cluster_config.Save()
def testCreateCommandContext_remote(self):
context = command_util.CommandContext('remotehost', 'remoteuser')
self.assertFalse(context.IsLocal())
self.assertEqual('remotehost', context.host)
self.assertEqual('remoteuser', context.user)
self.assertFalse(self.mock_ssh_context_creator.called)
self.mock_fabric_pkg.Connection.assert_called_once_with(
host='remotehost', user='******', config=mock.ANY)
def setUp(self):
super(LocalCommandContextTest, self).setUp()
self.subprocess_patcher = mock.patch(
'__main__.command_util.subprocess')
self.mock_subprocess_pkg = self.subprocess_patcher.start()
self.mock_process = mock.MagicMock(returncode=0)
self.mock_subprocess_pkg.Popen.return_value = self.mock_process
self.mock_process.communicate.return_value = ('stdout', 'stderr')
self.context = command_util.CommandContext()
def testClose_local(self):
context = command_util.CommandContext(
wrapped_context=self.wrapped_context)
self.assertTrue(context.IsLocal())
self.assertFalse(context._closed)
context.Close()
self.assertTrue(context._closed)
with self.assertRaises(command_util.CommandContextClosedError):
context.Run(['command'])
self.assertFalse(self.wrapped_context.close.called)
def testCreateCommandContext_local(self, get_hostname, get_user):
get_hostname.return_value = 'ahost'
get_user.return_value = 'auser'
context = command_util.CommandContext(socket.gethostname(),
getpass.getuser())
self.assertTrue(context.IsLocal())
self.assertEqual('ahost', context.host)
self.assertEqual('auser', context.user)
self.assertFalse(self.mock_fabric_pkg.Connection.called)
self.assertFalse(self.mock_ssh_context_creator.called)
def testCreateCommandContext_withSshConfig(self):
context = command_util.CommandContext(host='host1',
user='******',
ssh_config=ssh_util.SshConfig(
hostname='host1',
user='******'))
self.assertIsNotNone(context)
self.assertFalse(context.IsLocal())
self.assertFalse(self.mock_ssh_context_creator.called)
self.mock_fabric_pkg.Connection.assert_called_once_with(
host='host1', user='******', config=mock.ANY)
def testCopyFile_localSameFile(self, exists):
exists.return_value = True
context = command_util.CommandContext(
wrapped_context=self.wrapped_context)
self.assertTrue(context.IsLocal())
context.CopyFile('/local/file', '/local/file')
exists.assert_called_once_with('/local/file')
self.wrapped_context.run.assert_called_once_with('mkdir -p /local',
out_stream=mock.ANY,
err_stream=mock.ANY,
warn=True)
def testCreateCommandContext_useNativeSsh(self):
ssh_config = ssh_util.SshConfig(hostname='host1',
user='******',
use_native_ssh=True)
context = command_util.CommandContext(host='host1',
user='******',
ssh_config=ssh_config)
self.assertIsNotNone(context)
self.assertFalse(context.IsLocal())
self.mock_ssh_context_creator.assert_called_once_with(ssh_config)
self.assertFalse(self.mock_fabric_pkg.Connection.called)
def testFindGcloud(self):
context = command_util.CommandContext(
wrapped_context=self.wrapped_context)
self.assertFalse(context._tried_gcloud)
gcloud = context._FindGCloud()
self.assertEqual('gcloud', gcloud)
self.wrapped_context.assert_has_calls([
mock.call.run('which gcloud',
out_stream=mock.ANY,
err_stream=mock.ANY,
warn=True)
])
self.assertTrue(context._tried_gcloud)
def context(self):
"""Get remote context for the host."""
if not self._context:
try:
self.execution_state = 'Connect to host'
self._context = command_util.CommandContext(
self.name,
self.config.host_login_name,
ssh_config=self._ssh_config,
sudo_ssh_config=self._sudo_ssh_config)
except command_util.AuthenticationError as e:
logger.debug(e)
raise
return self._context
def RemoveNode(self, args):
"""Removes a node from an existing cluster.
Args:
args: an argparse.ArgumentParser object.
Raises:
ValueError: if a host or a token is missing.
"""
if not args.host:
raise ValueError('--host must be provided')
context = command_util.CommandContext(host=args.host,
user=args.ssh_user)
docker_context = command_util.DockerContext(context,
try_use_gcloud=False)
docker_context.Run(['swarm', 'leave', '--force'])
def _DownloadToolFromGCS(gcs_url, local_path):
"""Update local cli to remote cli with gcs."""
try:
cred = google_auth_util.GetGCloudCredential(
command_util.CommandContext())
gcs_client = gcs_file_util.CreateGCSClient(_MTT_PROJECT, cred)
blob = gcs_file_util.GetGCSBlob(gcs_client, gcs_url)
except gcs_file_util.GCSError as e:
logger.warning(e)
return None
local_md5hash = gcs_file_util.CalculateMd5Hash(local_path)
if six.ensure_text(blob.md5_hash) == local_md5hash:
logger.debug('Local is the same as remote, skip updating.')
return None
logger.info('There is a newer version %s on %s, updating.',
os.path.basename(local_path), gcs_url)
os.rename(local_path, gcs_file_util.CreateBackupFilePath(local_path))
blob.download_to_filename(local_path)
os.chmod(local_path, 0o770)
return local_path
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 testFindGcloud_noGcloud(self):
self.wrapped_context.run.side_effect = [
mock.MagicMock(return_code=1),
mock.MagicMock(return_code=1)
]
context = command_util.CommandContext(
wrapped_context=self.wrapped_context)
self.assertFalse(context._tried_gcloud)
gcloud = context._FindGCloud()
self.assertIsNone(gcloud)
self.wrapped_context.assert_has_calls([
mock.call.run('which gcloud',
out_stream=mock.ANY,
err_stream=mock.ANY,
warn=True),
mock.call.run('which /google/data/ro/teams/cloud-sdk/gcloud',
out_stream=mock.ANY,
err_stream=mock.ANY,
warn=True)
],
any_order=True)
self.assertTrue(context._tried_gcloud)
def setUp(self):
super(CommandContextTest, self).setUp()
self.fabric_patcher = mock.patch('__main__.command_util.fabric')
self.mock_fabric_pkg = self.fabric_patcher.start()
self.ssh_context_patcher = mock.patch('__main__.ssh_util.Context')
self.mock_ssh_context = mock.MagicMock()
self.mock_ssh_context_creator = self.ssh_context_patcher.start()
self.mock_ssh_context_creator.return_value = self.mock_ssh_context
self.wrapped_context = mock.MagicMock()
self.wrapped_context.run.return_value = mock.MagicMock(return_code=0)
self.sudo_wrapped_context = mock.MagicMock()
self.sudo_wrapped_context.sudo.return_value = mock.MagicMock(
return_code=0)
self.remote_context = command_util.CommandContext(
'remotehost',
'remoteuser',
ssh_config=ssh_util.SshConfig(hostname='remotehost',
user='******'),
sudo_ssh_config=ssh_util.SshConfig(hostname='remotehost',
user='******',
password='******'),
wrapped_context=self.wrapped_context,
sudo_wrapped_context=self.sudo_wrapped_context)
def AddNode(self, args):
"""Adds a node to an existing cluster.
Args:
args: an argparse.ArgumentParser object.
Raises:
ValueError: if a host or a token is missing.
"""
if not args.host:
raise ValueError('--host must be provided')
if not args.token:
raise ValueError('--token must be provided')
context = command_util.CommandContext(host=args.host,
user=args.ssh_user)
docker_context = command_util.DockerContext(context,
try_use_gcloud=False)
cluster_config = self._registry.GetConfig(args.name)
if args.host == cluster_config.manager_host:
raise ValueError('%s is already a manager node for %s cluster' %
(args.host, args.name))
docker_context.Run([
'swarm', 'join', '--token', args.token,
'%s:2377' % cluster_config.manager_host
])
def testCreateCommandContext_useNativeSsh_withSudo(self):
ssh_config = ssh_util.SshConfig(hostname='host1',
user='******',
use_native_ssh=True)
sudo_ssh_config = ssh_util.SshConfig(hostname='host1',
user='******',
use_native_ssh=True)
context = command_util.CommandContext(host='host1',
user='******',
ssh_config=ssh_config,
sudo_ssh_config=sudo_ssh_config)
self.assertIsNotNone(context)
self.assertFalse(context.IsLocal())
self.assertEqual(self.mock_ssh_context, context._sudo_wrapped_context)
self.assertEqual(self.mock_ssh_context, context._wrapped_context)
self.mock_ssh_context_creator.assert_has_calls([
mock.call(ssh_config),
mock.call(sudo_ssh_config),
mock.call().sudo('echo "Testing sudo access..."',
hide=True,
password=None)
])
self.assertFalse(self.mock_fabric_pkg.Connection.called)
def testCreateCommandContext(self):
context = command_util.CommandContext()
self.assertTrue(context.IsLocal())
self.assertFalse(self.mock_fabric_pkg.Connection.called)
self.assertFalse(self.mock_ssh_context_creator.called)
