def test_robust_begin_gcloud_ssh_pass(mocked_ssh, _):
"""Tests robust_begin_gcloud_ssh works as intended on google cloud."""
mocked_ssh.return_value = new_process.ProcessResult(0, None, False)
gcloud.robust_begin_gcloud_ssh(INSTANCE_NAME, ZONE)
mocked_ssh.assert_called_with('instance-a',
command='echo ping',
expect_zero=False,
zone='zone-a')
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_robust_begin_gcloud_ssh_fail(_, mocked_ssh):
"""Tests that ssh works as expected."""
with pytest.raises(Exception) as exception:
gcloud.robust_begin_gcloud_ssh(INSTANCE_NAME, ZONE)
assert mocked_ssh.call_count == 10
assert exception.value == 'Couldn\'t SSH to instance.'