def test_populate_job_yaml(job_spec_file, job, default_command_args):
environment = KubernetesJobEnvironment(
job_spec_file=job_spec_file, unique_job_name=True
)
job["spec"]["template"]["spec"]["containers"][0]["env"] = []
environment._job_spec = job
with set_temporary_config(
{
"cloud.graphql": "gql_test",
"cloud.auth_token": "auth_test",
"logging.extra_loggers": "['test_logger']",
}
):
with prefect.context(flow_run_id="id_test", namespace="namespace_test"):
yaml_obj = environment._populate_run_time_job_spec_details(
docker_name="test1/test2:test3",
)
assert "prefect-dask-job-" in yaml_obj["metadata"]["name"]
assert len(yaml_obj["metadata"]["name"]) == 25
assert (
yaml_obj["metadata"]["labels"]["prefect.io/identifier"]
== environment.identifier_label
)
assert yaml_obj["metadata"]["labels"]["prefect.io/flow_run_id"] == "id_test"
assert (
yaml_obj["spec"]["template"]["metadata"]["labels"]["prefect.io/identifier"]
== environment.identifier_label
)
env = yaml_obj["spec"]["template"]["spec"]["containers"][0]["env"]
assert env[0]["value"] == "gql_test"
assert env[1]["value"] == "auth_test"
assert env[2]["value"] == "id_test"
assert env[3]["value"] == "namespace_test"
assert env[4]["value"] == "test1/test2:test3"
assert env[9]["value"] == "['test_logger']"
assert (
yaml_obj["spec"]["template"]["spec"]["containers"][0]["image"]
== "test1/test2:test3"
)
assert yaml_obj["spec"]["template"]["spec"]["containers"][0]["command"] == [
"/bin/sh",
"-c",
]
assert (
yaml_obj["spec"]["template"]["spec"]["containers"][0]["args"]
== default_command_args
)
def test_populate_job_yaml_no_defaults(job_spec_file, job):
environment = KubernetesJobEnvironment(job_spec_file=job_spec_file)
# only command and args are set on the container when the instance
# is initialized
job["spec"]["template"]["spec"]["containers"][0] = {
"command": ["/bin/sh", "-c"],
"args": default_command_args,
}
del job["metadata"]
del job["spec"]["template"]["metadata"]
environment._job_spec = job
with set_temporary_config(
{"cloud.graphql": "gql_test", "cloud.auth_token": "auth_test"}
):
with prefect.context(flow_run_id="id_test", namespace="namespace_test"):
yaml_obj = environment._populate_run_time_job_spec_details(
docker_name="test1/test2:test3",
)
assert (
yaml_obj["metadata"]["labels"]["prefect.io/identifier"]
== environment.identifier_label
)
assert yaml_obj["metadata"]["labels"]["prefect.io/flow_run_id"] == "id_test"
assert (
yaml_obj["spec"]["template"]["metadata"]["labels"]["prefect.io/identifier"]
== environment.identifier_label
)
env = yaml_obj["spec"]["template"]["spec"]["containers"][0]["env"]
assert env[0]["value"] == "gql_test"
assert env[1]["value"] == "auth_test"
assert env[2]["value"] == "id_test"
assert env[3]["value"] == "namespace_test"
assert env[4]["value"] == "test1/test2:test3"
assert env[9]["value"] == "[]"
assert (
yaml_obj["spec"]["template"]["spec"]["containers"][0]["image"]
== "test1/test2:test3"
)
def test_populate_job_yaml_multiple_containers(
job_spec_file, job, default_command_args
):
environment = KubernetesJobEnvironment(job_spec_file=job_spec_file)
# Generate yaml object with multiple containers
job["spec"]["template"]["spec"]["containers"][0]["env"] = []
job["spec"]["template"]["spec"]["containers"].append(
copy.deepcopy(job["spec"]["template"]["spec"]["containers"][0])
)
job["spec"]["template"]["spec"]["containers"][1]["env"] = []
job["spec"]["template"]["spec"]["containers"][1]["args"] = "echo 'other command'"
environment._job_spec = job
with set_temporary_config(
{
"cloud.graphql": "gql_test",
"cloud.auth_token": "auth_test",
"logging.extra_loggers": "['test_logger']",
}
):
with prefect.context(flow_run_id="id_test", namespace="namespace_test"):
yaml_obj = environment._populate_run_time_job_spec_details(
docker_name="test1/test2:test3",
)
assert (
yaml_obj["metadata"]["labels"]["prefect.io/identifier"]
== environment.identifier_label
)
assert yaml_obj["metadata"]["labels"]["prefect.io/flow_run_id"] == "id_test"
assert (
yaml_obj["spec"]["template"]["metadata"]["labels"]["prefect.io/identifier"]
== environment.identifier_label
)
# Assert First Container
env = yaml_obj["spec"]["template"]["spec"]["containers"][0]["env"]
assert env[0]["value"] == "gql_test"
assert env[1]["value"] == "auth_test"
assert env[2]["value"] == "id_test"
assert env[3]["value"] == "namespace_test"
assert env[4]["value"] == "test1/test2:test3"
assert env[9]["value"] == "['test_logger']"
assert (
yaml_obj["spec"]["template"]["spec"]["containers"][0]["image"]
== "test1/test2:test3"
)
assert yaml_obj["spec"]["template"]["spec"]["containers"][0]["command"] == [
"/bin/sh",
"-c",
]
assert (
yaml_obj["spec"]["template"]["spec"]["containers"][0]["args"]
== default_command_args
)
# Assert Second Container
env = yaml_obj["spec"]["template"]["spec"]["containers"][1]["env"]
assert env[0]["value"] == "gql_test"
assert env[1]["value"] == "auth_test"
assert env[2]["value"] == "id_test"
assert env[3]["value"] == "namespace_test"
assert env[4]["value"] == "test1/test2:test3"
assert env[9]["value"] == "['test_logger']"
assert (
yaml_obj["spec"]["template"]["spec"]["containers"][1]["image"]
!= "test1/test2:test3"
)
assert (
yaml_obj["spec"]["template"]["spec"]["containers"][1]["args"]
!= default_command_args
)
