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()