def testHostContext_withSSHInfo(self):
ssh_config = ssh_util.SshConfig(user='******',
hostname='host1',
password='******',
ssh_key='/ssh_key')
sudo_ssh_config = ssh_util.SshConfig(user='******',
hostname='host1',
password='******',
ssh_key='/ssh_key')
host = host_util.Host(self.host_config1,
ssh_config=ssh_config,
sudo_ssh_config=sudo_ssh_config)
self.assertIsNotNone(host.context)
self.mock_create_context.assert_called_once_with(
'host1',
'user1',
ssh_config=ssh_config,
sudo_ssh_config=sudo_ssh_config)
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 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 testCreateFabricContext_withPassword(self):
context = command_util._CreateFabricContext(
ssh_util.SshConfig(user='******', hostname='host1', password='******'))
self.assertIsNotNone(context)
self.mock_fabric_pkg.Config.assert_called_once_with(
system_ssh_path='~/.ssh/config')
self.mock_fabric_pkg.Connection.assert_called_once_with(
host='host1',
user='******',
config=mock.ANY,
connect_kwargs={'password': '******'})
self.mock_fabric_pkg.Connection().open.assert_called_once_with()
def testExecute_askPass(self, mock_build_lab_config_pool, mock_getpass):
"""Test execute on multiple hosts with password."""
mock_getpass.side_effect = ['login_pswd', 'sudo_pswd']
mock_build_lab_config_pool.return_value = self.mock_lab_config_pool
self.mock_lab_config_pool.GetHostConfigs.side_effect = [[
self.host_config1, self.host_config2
]]
args_dict = self.default_args.copy()
args_dict.update(
ask_login_password=True,
ask_sudo_password=True,
sudo_user='******',
service_account_json_key_path='path/to/key',
)
args = mock.MagicMock(**args_dict)
host_util.Execute(args)
self.mock_lab_config_pool.GetHostConfigs.assert_called_once_with()
self.mock_create_context.assert_has_calls([
mock.call('host1',
'user1',
ssh_config=ssh_util.SshConfig(user='******',
hostname='host1',
password='******'),
sudo_ssh_config=ssh_util.SshConfig(
user='******',
hostname='host1',
password='******')),
mock.call('host2',
'user1',
ssh_config=ssh_util.SshConfig(user='******',
hostname='host2',
password='******'),
sudo_ssh_config=ssh_util.SshConfig(user='******',
hostname='host2',
password='******'))
])
self.assertSameElements(['host1', 'host2'],
self.mock_func_calls.keys())
def __init__(self, host=None, user=None,
ssh_config=None, sudo_ssh_config=None,
wrapped_context=None, sudo_wrapped_context=None):
"""Create context that will be used to run command with.
Args:
host: a host name.
user: username of the host.
ssh_config: ssh_util.SshConfig for login to the host.
sudo_ssh_config: ssh_util.SshConfig for login to the host as sudo.
wrapped_context: invoke context, used for test.
sudo_wrapped_context: invoke context, used for sudo test.
Raises:
FabricNotFoundError: fabric is not installed.
"""
self._host = host or socket.gethostname()
self._user = user or getpass.getuser()
self._is_local = False
self._tried_gcloud = False
self._gcloud = None
if socket.gethostname() == self._host or socket.getfqdn() == self._host:
self._is_local = True
self._ssh_config = ssh_config or ssh_util.SshConfig(
user=user, hostname=host)
self._sudo_ssh_config = sudo_ssh_config
self._closed = False
# TODO: We should make the subprocess's output configurable by
# parent process.
self._out_stream = HostCommandOutStream(
None if self._is_local else host, 'stdout')
# We also output stderr to DEBUG level, if the command failed and
# raise on failure, the raised exception will have the information.
self._err_stream = HostCommandOutStream(
None if self._is_local else host, 'stderr')
self._wrapped_context = wrapped_context
self._sudo_wrapped_context = sudo_wrapped_context
if self._wrapped_context:
return
if self._is_local:
if invoke:
# Use invoke context to running locally.
self._wrapped_context = invoke.Context()
# Else use _RunDirectly instead of using invoke.context.
return
if self._ssh_config.use_native_ssh or not fabric:
if not self._ssh_config.use_native_ssh and not fabric:
logger.debug('No fabric package installed, using native ssh.')
self._CreateNativeSshContexts()
return
self._CreateFabricContexts()
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 testCreateFabricContext_withSshKey(self):
context = command_util._CreateFabricContext(
ssh_util.SshConfig(user='******',
hostname='host1',
ssh_key='/sshkey'))
self.assertIsNotNone(context)
self.mock_fabric_pkg.Config.assert_called_once_with(
system_ssh_path='~/.ssh/config')
self.mock_fabric_pkg.Connection.assert_called_once_with(
host='host1',
user='******',
config=mock.ANY,
connect_kwargs={'key_filename': '/sshkey'})
self.mock_fabric_pkg.Connection().open.assert_called_once_with()
def setUp(self):
super(HostUtilTest, self).setUp()
self.mock_context = mock.MagicMock()
self.context_patcher = mock.patch(
'__main__.host_util.command_util.CommandContext',
return_value=self.mock_context)
self.mock_create_context = self.context_patcher.start()
self.host_config1 = host_util.lab_config.CreateHostConfig(
cluster_name='cluster1',
hostname='host1',
host_login_name='user1',
docker_image='image1')
self.ssh_config1 = ssh_util.SshConfig(user='******', hostname='host1')
self.host_config2 = host_util.lab_config.CreateHostConfig(
cluster_name='cluster2',
hostname='host2',
host_login_name='user1',
docker_image='image1')
self.ssh_config2 = ssh_util.SshConfig(user='******', hostname='host2')
self.host_config3 = host_util.lab_config.CreateHostConfig(
cluster_name='cluster3', hostname='host3', host_login_name='user1')
self.ssh_config3 = ssh_util.SshConfig(user='******', hostname='host3')
self.mock_lab_config_pool = mock.MagicMock()
self.mock_func_calls = collections.OrderedDict()
self.mock_func_exceptions = {}
self.default_args = {
'lab_config_path': 'lab_config.yaml',
'hosts_or_clusters': [],
'parallel': False,
'host_func': self._MockFunc,
'ssh_key': None,
'ask_login_password': False,
'ask_sudo_password': False,
'sudo_user': None,
'use_native_ssh': False,
'ssh_arg': None,
}
def testExecute_nativeSsh(self, mock_build_lab_config_pool):
"""Test execute."""
mock_build_lab_config_pool.return_value = self.mock_lab_config_pool
self.mock_lab_config_pool.GetHostConfigs.side_effect = [[
self.host_config1, self.host_config2
]]
args_dict = self.default_args.copy()
args_dict.update(
use_native_ssh=True,
ssh_arg='-o op1=v1 -o op2=v2',
service_account_json_key_path='path/to/key',
)
args = mock.MagicMock(**args_dict)
host_util.Execute(args)
self.mock_lab_config_pool.GetHostConfigs.assert_called_once_with()
self.mock_create_context.assert_has_calls([
mock.call('host1',
'user1',
ssh_config=ssh_util.SshConfig(
user='******',
hostname='host1',
ssh_args='-o op1=v1 -o op2=v2',
use_native_ssh=True),
sudo_ssh_config=None),
mock.call('host2',
'user1',
ssh_config=ssh_util.SshConfig(
user='******',
hostname='host2',
ssh_args='-o op1=v1 -o op2=v2',
use_native_ssh=True),
sudo_ssh_config=None)
])
self.assertSameElements(['host1', 'host2'],
self.mock_func_calls.keys())
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 testRun(self):
"""Test run."""
c = ssh_util.Context(ssh_util.SshConfig(user='******',
hostname='ahost'))
res = c.run('/tmp/mtt start /tmp/lab.yaml')
self.mock_subprocess_pkg.Popen.assert_called_once_with(
[
'ssh', '-o', 'User=auser', 'ahost',
'/bin/sh -c \'/tmp/mtt start /tmp/lab.yaml\''
],
stdin=self.mock_subprocess_pkg.DEVNULL,
stdout=self.mock_subprocess_pkg.PIPE,
stderr=self.mock_subprocess_pkg.PIPE,
universal_newlines=True)
self.assertEqual(0, res.return_code)
def testPut(self):
"""Test put."""
c = ssh_util.Context(ssh_util.SshConfig(user='******',
hostname='ahost'))
c.put('/path/to/local/file', '/path/to/remote/file')
self.mock_subprocess_pkg.Popen.assert_called_once_with(
[
'rsync', '/path/to/local/file',
'auser@ahost:/path/to/remote/file'
],
stdin=self.mock_subprocess_pkg.DEVNULL,
stdout=self.mock_subprocess_pkg.PIPE,
stderr=self.mock_subprocess_pkg.PIPE,
universal_newlines=True)
self.assertTrue(self.mock_process.communicate.called)
def testRun_withSshArgs_notTokenized(self):
"""Test run with ssh args."""
c = ssh_util.Context(
ssh_util.SshConfig(user='******',
hostname='ahost',
ssh_args=('-o StrictHostKeyChecking=no '
'-o UserKnownHostsFile=/dev/null')))
res = c.run('/tmp/mtt start /tmp/lab.yaml')
self.mock_subprocess_pkg.Popen.assert_called_once_with(
[
'ssh', '-o', 'StrictHostKeyChecking=no', '-o',
'UserKnownHostsFile=/dev/null', '-o', 'User=auser', 'ahost',
'/bin/sh -c \'/tmp/mtt start /tmp/lab.yaml\''
],
stdin=self.mock_subprocess_pkg.DEVNULL,
stdout=self.mock_subprocess_pkg.PIPE,
stderr=self.mock_subprocess_pkg.PIPE,
universal_newlines=True)
self.assertEqual(0, res.return_code)
def testPut_withSshArgs(self):
"""Test put with ssh args."""
c = ssh_util.Context(
ssh_util.SshConfig(user='******',
hostname='ahost',
ssh_args=('-o StrictHostKeyChecking=no '
'-o UserKnownHostsFile=/dev/null')))
c.put('/path/to/local/file', '/path/to/remote/file')
self.mock_subprocess_pkg.Popen.assert_called_once_with(
[
'rsync', '-e',
'ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null',
'/path/to/local/file', 'auser@ahost:/path/to/remote/file'
],
stdin=self.mock_subprocess_pkg.DEVNULL,
stdout=self.mock_subprocess_pkg.PIPE,
stderr=self.mock_subprocess_pkg.PIPE,
universal_newlines=True)
self.assertTrue(self.mock_process.communicate.called)
def Execute(args, lab_config_pool=None):
"""Execute a command on hosts."""
lab_config_pool = lab_config_pool or BuildLabConfigPool(
args.lab_config_path, args.service_account_json_key_path)
host_configs = _GetHostConfigs(lab_config_pool, args.hosts_or_clusters)
if not host_configs:
logger.warning('No host configured in %s for %s.',
args.lab_config_path,
args.hosts_or_clusters)
return
login_password = None
if args.ask_login_password:
login_password = getpass.getpass('Enter the login password:'******'Enter the sudo password:'******'Using native ssh instead of fabric.')
if args.ssh_arg:
logger.debug('Use ssh arg: %s', args.ssh_arg)
host_configs.sort(key=lambda host_config: host_config.hostname)
hosts = []
for host_config in host_configs:
if (args.service_account_json_key_path and
not host_config.service_account_json_key_path):
host_config = host_config.SetServiceAccountJsonKeyPath(
args.service_account_json_key_path)
ssh_args = args.ssh_arg or host_config.ssh_arg
# only native ssh support ssh_args
use_native_ssh = bool(args.use_native_ssh or ssh_args)
ssh_config = ssh_util.SshConfig(
hostname=host_config.hostname,
user=host_config.host_login_name,
password=login_password,
ssh_args=ssh_args,
ssh_key=args.ssh_key,
use_native_ssh=use_native_ssh)
sudo_ssh_config = None
if args.sudo_user or sudo_password:
sudo_ssh_config = ssh_util.SshConfig(
hostname=host_config.hostname,
user=args.sudo_user or host_config.host_login_name,
password=sudo_password or login_password,
ssh_args=ssh_args,
ssh_key=args.ssh_key,
use_native_ssh=use_native_ssh)
hosts.append(
Host(host_config,
ssh_config=ssh_config,
sudo_ssh_config=sudo_ssh_config))
execution_state_printer = ExecutionStatePrinter(hosts)
try:
f = _WrapFuncForSetHost(args.host_func)
if args.parallel:
_ParallelExecute(f, args, hosts, execution_state_printer)
else:
_SequentialExecute(
f, args, hosts,
exit_on_error=args.exit_on_error,
execution_state_printer=execution_state_printer)
except KeyboardInterrupt:
logger.info('Receive KeyboardInterrupt.')
finally:
for host in hosts:
try:
if host.context:
logger.debug('Closing %s.', host.name)
host.context.Close()
except Exception as e: logger.error('Failed to close %s: %s.', host.name, e)
logger.info('Finished executing "%s".', args.host_func.__name__)
execution_state_printer.PrintResult()
