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)