def test_ssh(): """Tests that ssh works as expected.""" with test_utils.mock_popen_ctx_mgr() as mocked_popen: gcloud.ssh(INSTANCE_NAME) assert mocked_popen.commands == [[ 'gcloud', 'beta', 'compute', 'ssh', INSTANCE_NAME ]]
def start(self): """Start the experiment on the dispatcher.""" # TODO(metzman): Replace this workflow with a startup script so we don't # need to SSH into the dispatcher. self.process.join() # Wait for dispatcher instance. # Check that we can SSH into the instance. gcloud.robust_begin_gcloud_ssh(self.instance_name, self.config['cloud_compute_zone']) base_docker_tag = experiment_utils.get_base_docker_tag( self.config['cloud_project']) cloud_sql_instance_connection_name = ( self.config['cloud_sql_instance_connection_name']) command = ( 'echo 0 | sudo tee /proc/sys/kernel/yama/ptrace_scope && ' 'docker run --rm ' '-e INSTANCE_NAME="{instance_name}" ' '-e EXPERIMENT="{experiment}" ' '-e CLOUD_PROJECT="{cloud_project}" ' '-e CLOUD_EXPERIMENT_BUCKET="{cloud_experiment_bucket}" ' '-e POSTGRES_PASSWORD="******" ' '-e CLOUD_SQL_INSTANCE_CONNECTION_NAME=' '"{cloud_sql_instance_connection_name}" ' '--cap-add=SYS_PTRACE --cap-add=SYS_NICE ' '-v /var/run/docker.sock:/var/run/docker.sock ' '--name=dispatcher-container ' '{base_docker_tag}/dispatcher-image ' '/work/startup-dispatcher.sh' ).format( instance_name=self.instance_name, postgres_password=os.environ['POSTGRES_PASSWORD'], experiment=self.config['experiment'], # TODO(metzman): Create a function that sets env vars based on # the contents of a dictionary, and use it instead of hardcoding # the configs we use. cloud_project=self.config['cloud_project'], cloud_experiment_bucket=self.config['cloud_experiment_bucket'], cloud_sql_instance_connection_name=( cloud_sql_instance_connection_name), base_docker_tag=base_docker_tag, ) return gcloud.ssh(self.instance_name, command=command, zone=self.config['cloud_compute_zone'])
def test_ssh_optional_arg(kwargs_for_ssh, expected_argument): """Tests that ssh works as expected when given an optional argument.""" with test_utils.mock_popen_ctx_mgr() as mocked_popen: gcloud.ssh(INSTANCE_NAME, **kwargs_for_ssh) assert expected_argument in mocked_popen.commands[0]