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]