def test_job_actually_runs_script(self):
with CloudSdkIntegrationStub() as gcloud:
job = clash.Job(gcloud=gcloud, job_config=TEST_JOB_CONFIG)
job.run("echo hello")
assert b"hello\n" in gcloud.instances[0].logs()
def test_attaching_raises_exception_after_timeout(self):
self.gcloud.get_subscriber().pull.return_value.received_messages = []
job = clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud)
job.run(args=[])
with pytest.raises(TimeoutError) as e_info:
result = job.attach(timeout_seconds=1)
def test_running_a_job_creates_a_topic_path(self):
job = clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud)
job.run(args=[])
self.gcloud.get_publisher().topic_path.assert_called_with(
TEST_JOB_CONFIG["project_id"], job.name)
def test_job_runs_script_from_file(self):
with CloudSdkIntegrationStub() as gcloud:
job = clash.Job(gcloud=gcloud, job_config=TEST_JOB_CONFIG)
job.run_file("tests/script.sh")
assert b"hello\nworld\n" in gcloud.instances[0].logs()
def test_passing_gcs_target_without_artifacts_shows_error(self):
with CloudSdkIntegrationStub() as gcloud:
job = clash.Job(gcloud=gcloud, job_config=TEST_JOB_CONFIG)
job.run("", gcs_target={"/tmp/artifacts": "mybucket"})
assert b"No artifacts found in /tmp/artifacts" in gcloud.instances[0].logs()
def test_running_a_job_creates_a_pubsub_topic(self, mock_uuid_call):
mock_uuid_call.return_value = 1234
job = clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud)
job.run("")
self.gcloud.get_publisher().create_topic.assert_called_with(
f"{TEST_JOB_CONFIG['project_id']}/clash-job-1234")
def test_attaching_succeeds_if_there_is_a_running_job_and_a_message(self):
message = MagicMock()
message.message = MagicMock(data='{"status": 0}')
self.gcloud.get_subscriber().pull.return_value.received_messages = [message]
job = clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud)
job.run("")
job.attach() # throws no exception
def test_running_a_job_creates_an_instance_template(self):
job = clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud)
job.run(args=[])
self.gcloud.get_compute_client(
).instanceTemplates.return_value.insert.return_value.execute.assert_called(
)
def test_running_a_job_creates_a_managed_instance_group(self):
job = clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud)
job.run(args=[])
self.gcloud.get_compute_client(
).instanceGroupManagers.return_value.insert.return_value.execute.assert_called(
)
def test_attaching_returns_status_code(self):
message = MagicMock()
message.message = MagicMock(data='{"status": 127}')
self.gcloud.get_subscriber().pull.return_value.received_messages = [message]
job = clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud)
job.run("")
result = job.attach()
assert result["status"] == 127
def test_passing_gcs_mount_invokes_gcsfuse(self):
with CloudSdkIntegrationStub() as gcloud:
job = clash.Job(gcloud=gcloud, job_config=TEST_JOB_CONFIG)
job.run("", gcs_mounts={"mybucket": "/mnt/static"})
assert (
b"gcsfuse.--implicit-dirs.mybucket./mnt/static"
in gcloud.instances[0].logs()
)
def test_job_uses_given_env_vars(self):
with CloudSdkIntegrationStub() as gcloud:
script = """
echo "$MESSAGE"
"""
job = clash.Job(gcloud=gcloud, job_config=TEST_JOB_CONFIG)
job.run(script, env_vars={"MESSAGE": "foobar"})
assert b"foobar\n" in gcloud.instances[0].logs()
def test_job_shutdowns_machine_eventually(self):
with CloudSdkIntegrationStub() as gcloud:
job = clash.Job(gcloud=gcloud, job_config=TEST_JOB_CONFIG)
job.run("echo hello")
assert (
b"gcloud.compute.instance-groups.managed.delete"
in gcloud.instances[0].logs()
)
def test_running_a_job_creates_a_pubsub_topic(self, mock_uuid_call):
mock_uuid_call.return_value = 1234
job = clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud)
job.run(args=[])
self.gcloud.get_publisher().create_topic.assert_called_with(
f"{TEST_JOB_CONFIG['project_id']}/clash-job-1234",
message_storage_policy=MessageStoragePolicy(
allowed_persistence_regions=["europe-west1"]),
)
def test_removes_topic_if_job_creation_failed(self):
self.gcloud.get_compute_client(
).instanceGroupManagers.return_value.insert.return_value.execute.side_effect = Exception(
"Failure!")
job = clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud)
with pytest.raises(Exception) as e_info:
job.run(args=[])
self.gcloud.get_publisher().delete_topic.assert_called_with(
f"{TEST_JOB_CONFIG['project_id']}/{job.name}")
def test_on_finish_does_not_run_callback_when_job_is_still_running(self):
job = clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud)
result = {"called": False}
job.run(args=[])
def callback(status_code):
result["called"] = True
job.on_finish(callback)
assert not result["called"]
def test_on_finish_acknowledges_message(self):
message = MagicMock()
message.data = '{ "status": 0 }'
self.gcloud.get_subscriber().subscribe.side_effect = (
lambda path, callback: callback(message))
job = clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud)
job.run(args=[])
job.on_finish(lambda status_code: None)
message.ack.assert_called()
def test_job_sends_pubpub_message_on_failure(self, mock_uuid_call):
mock_uuid_call.return_value = 123
with CloudSdkIntegrationStub() as gcloud:
job = clash.Job(gcloud=gcloud, job_config=TEST_JOB_CONFIG)
job.run("exit 1")
assert (
b'gcloud.pubsub.topics.publish.clash-job-123.--message={"status": 1}'
in gcloud.instances[0].logs()
)
def test_deletes_instance_template_after_job_is_complete(self):
self.gcloud.get_compute_client(
).instanceGroups.return_value.list.return_value.execute.return_value = {
"items": [{
"name": "anothergroup"
}]
}
with clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud) as job:
job.run(args=[])
self.gcloud.get_compute_client(
).instanceTemplates.return_value.delete.return_value.execute.assert_called(
)
def test_job_runs_multiline_script(self):
with CloudSdkIntegrationStub() as gcloud:
script = """
echo 'hello'
the_world_is_flat=true
if [ "$the_world_is_flat" = true ] ; then
echo 'world'
fi
"""
job = clash.Job(gcloud=gcloud, job_config=TEST_JOB_CONFIG)
job.run(script)
assert b"hello\nworld\n" in gcloud.instances[0].logs()
def test_attaching_pulls_message(self):
message = MagicMock()
message.message = MagicMock(data='{"status": 0}')
self.gcloud.get_subscriber().pull.return_value.received_messages = [message]
self.gcloud.get_subscriber().subscription_path.side_effect = (
lambda x, y: "mysubscription"
)
job = clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud)
job.run("")
result = job.attach()
self.gcloud.get_subscriber().pull.assert_called_with(
"mysubscription", max_messages=1, return_immediately=False, timeout=30
)
def test_on_finish_runs_callback_when_job_is_complete(self):
message = MagicMock()
message.data = '{ "status": 0 }'
self.gcloud.get_subscriber().subscribe.side_effect = (
lambda path, callback: callback(message))
job = clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud)
result = {"status": -1}
job.run(args=[])
def callback(status_code):
result["status"] = status_code
job.on_finish(callback)
assert result["status"] == 0
def test_attaching_acknowledges_messages(self):
message = MagicMock(ack_id=42)
message.message = MagicMock(data='{"status": 0}')
self.gcloud.get_subscriber().pull.return_value.received_messages = [
message
]
self.gcloud.get_subscriber().subscription_path.side_effect = (
lambda x, y: "mysubscription")
job = clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud)
job.run(args=[])
result = job.attach()
self.gcloud.get_subscriber().acknowledge.assert_called_with(
"mysubscription", [42])
def test_running_a_job_creates_a_pubsub_subscription_for_status_updates(self):
message = MagicMock()
message.message = MagicMock(data='{"status": 0}')
self.gcloud.get_subscriber().pull.return_value.received_messages = [message]
self.gcloud.get_publisher().topic_path.side_effect = lambda x, y: "mytopic"
self.gcloud.get_subscriber().topic_path.side_effect = lambda x, y: "mytopic"
self.gcloud.get_subscriber().subscription_path.side_effect = (
lambda x, y: "mysubscription"
)
job = clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud)
job.run("")
self.gcloud.get_subscriber().create_subscription.assert_called_with(
"mysubscription", "mytopic"
)
def __init__(self,
job_config,
cmd=None,
cmd_file=None,
name_prefix=None,
env_vars={},
gcs_target={},
gcs_mounts={},
*args,
**kwargs):
self.job = clash.Job(job_config=job_config, name_prefix=name_prefix)
self.cmd = cmd
self.cmd_file = cmd_file
self.env_vars = env_vars
self.gcs_target = gcs_target
self.gcs_mounts = gcs_mounts
super(ComputeEngineJobOperator, self).__init__(*args, **kwargs)
def test_passing_gcs_target_invokes_gsutil(self):
with CloudSdkIntegrationStub() as gcloud:
script = """
touch /tmp/artifacts/foo
touch /tmp/models/bar
"""
job = clash.Job(gcloud=gcloud, job_config=TEST_JOB_CONFIG)
job.run(
script,
gcs_target={
"/tmp/artifacts": "mybucket",
"/tmp/models": "modelsbucket",
},
)
assert (
b"gsutil.cp.-r./tmp/artifacts/foo.gs://mybucket"
in gcloud.instances[0].logs()
)
assert (
b"gsutil.cp.-r./tmp/models/bar.gs://modelsbucket"
in gcloud.instances[0].logs()
)
def test_creates_job_with_name_prefix(self, mock_uuid_call):
mock_uuid_call.return_value = 1234
job = clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud, name_prefix="foo")
assert "foo-clash-job-1234" == job.name
def test_on_finish_fails_if_job_is_not_running(self):
job = clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud)
with pytest.raises(ValueError) as e_info:
job.on_finish(lambda status_code: None)
def test_creates_job(self, mock_uuid_call):
mock_uuid_call.return_value = 1234
job = clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud)
assert "clash-job-1234" == job.name
def test_attaching_fails_if_there_is_not_a_running_job(self):
job = clash.Job(TEST_JOB_CONFIG, gcloud=self.gcloud)
with pytest.raises(ValueError) as e_info:
job.attach()