def test_k8s_agent_replace_yaml_respects_multiple_image_secrets(monkeypatch, cloud_api):
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
monkeypatch.setenv("IMAGE_PULL_SECRETS", "some-secret,other-secret")
monkeypatch.setenv("IMAGE_PULL_POLICY", "custom_policy")
flow_run = GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": Docker(
registry_url="test", image_name="name", image_tag="tag"
).serialize(),
"environment": LocalEnvironment().serialize(),
"id": "new_id",
"core_version": "0.13.0",
}
),
"id": "id",
}
)
with set_temporary_config(
{"cloud.agent.auth_token": "token", "logging.log_to_cloud": True}
):
agent = KubernetesAgent(env_vars=dict(AUTH_THING="foo", PKG_SETTING="bar"))
job = agent.generate_job_spec_from_environment(flow_run, image="test/name:tag")
expected_secrets = [{"name": "some-secret"}, {"name": "other-secret"}]
assert job["spec"]["template"]["spec"]["imagePullSecrets"] == expected_secrets
def test_k8s_agent_generate_deployment_yaml(monkeypatch, cloud_api):
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
agent = KubernetesAgent()
deployment = agent.generate_deployment_yaml(
token="test_token",
api="test_api",
namespace="test_namespace",
resource_manager_enabled=True,
backend="backend-test",
)
deployment = yaml.safe_load(deployment)
agent_env = deployment["spec"]["template"]["spec"]["containers"][0]["env"]
resource_manager_env = deployment["spec"]["template"]["spec"][
"containers"][1]["env"]
assert agent_env[0]["value"] == "test_token"
assert agent_env[1]["value"] == "test_api"
assert agent_env[2]["value"] == "test_namespace"
assert agent_env[11]["value"] == "backend-test"
assert resource_manager_env[0]["value"] == "test_token"
assert resource_manager_env[1]["value"] == "test_api"
assert resource_manager_env[3]["value"] == "test_namespace"
def test_generate_job_spec_image_pull_secrets_empty_string_in_runconfig(
self, tmpdir):
"""Regression test for issue #5001."""
run_config = KubernetesRun(image_pull_secrets="")
agent = KubernetesAgent(namespace="testing")
job = agent.generate_job_spec(self.build_flow_run(run_config))
assert "imagePullSecrets" not in job["spec"]["template"]["spec"]
def test_k8s_agent_generate_deployment_yaml_contains_resources(
monkeypatch, cloud_api):
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
agent = KubernetesAgent()
deployment = agent.generate_deployment_yaml(
token="test_token",
api="test_api",
namespace="test_namespace",
mem_request="mr",
mem_limit="ml",
cpu_request="cr",
cpu_limit="cl",
image_pull_policy="custom_policy",
service_account_name="svc",
)
deployment = yaml.safe_load(deployment)
env = deployment["spec"]["template"]["spec"]["containers"][0]["env"]
assert env[5]["value"] == "mr"
assert env[6]["value"] == "ml"
assert env[7]["value"] == "cr"
assert env[8]["value"] == "cl"
assert env[9]["value"] == "custom_policy"
assert env[10]["value"] == "svc"
def test_k8s_agent_manage_jobs_pass(monkeypatch, cloud_api):
job_mock = MagicMock()
job_mock.metadata.labels = {
"prefect.io/identifier": "id",
"prefect.io/flow_run_id": "fr",
}
job_mock.metadata.name = "my_job"
batch_client = MagicMock()
list_job = MagicMock()
list_job.metadata._continue = 0
list_job.items = [job_mock]
batch_client.list_namespaced_job.return_value = list_job
monkeypatch.setattr("kubernetes.client.BatchV1Api",
MagicMock(return_value=batch_client))
pod = MagicMock()
pod.metadata.name = "pod_name"
core_client = MagicMock()
list_pods = MagicMock()
list_pods.items = [pod]
core_client.list_namespaced_pod.return_value = list_pods
monkeypatch.setattr("kubernetes.client.CoreV1Api",
MagicMock(return_value=core_client))
agent = KubernetesAgent()
agent.heartbeat()
def test_k8s_agent_replace_yaml_no_pull_secrets(monkeypatch, cloud_api):
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
LocalEnvironment().serialize(),
"id":
"id",
"core_version":
"0.13.0",
}),
"id":
"id",
})
agent = KubernetesAgent()
job = agent.generate_job_spec_from_environment(flow_run,
image="test/name:tag")
assert not job["spec"]["template"]["spec"].get("imagePullSecrets", None)
def test_k8s_agent_deploy_flow_raises(monkeypatch, cloud_api):
batch_client = MagicMock()
monkeypatch.setattr("kubernetes.client.BatchV1Api",
MagicMock(return_value=batch_client))
core_client = MagicMock()
core_client.list_namespaced_pod.return_value = MagicMock(items=[])
monkeypatch.setattr("kubernetes.client.CoreV1Api",
MagicMock(return_value=core_client))
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
agent = KubernetesAgent()
with pytest.raises(ValueError):
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"storage": Local().serialize(),
"id": "id",
"environment": LocalEnvironment().serialize(),
"core_version": "0.13.0",
}),
"id":
"id",
}))
assert not agent.batch_client.create_namespaced_job.called
def test_k8s_agent_replace_yaml_responds_to_logging_config(
monkeypatch, cloud_api, flag):
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
LocalEnvironment().serialize(),
"id":
"new_id",
"core_version":
"0.13.0",
}),
"id":
"id",
"name":
"name",
})
agent = KubernetesAgent(no_cloud_logs=flag)
job = agent.generate_job_spec_from_environment(flow_run,
image="test/name:tag")
env = job["spec"]["template"]["spec"]["containers"][0]["env"]
assert env[6]["value"] == str(not flag).lower()
def test_k8s_agent_manage_jobs_delete_jobs(monkeypatch, cloud_api):
job_mock = MagicMock()
job_mock.metadata.labels = {
"prefect.io/identifier": "id",
"prefect.io/flow_run_id": "fr",
}
job_mock.metadata.name = "my_job"
job_mock.status.failed = True
job_mock.status.succeeded = True
batch_client = MagicMock()
list_job = MagicMock()
list_job.metadata._continue = 0
list_job.items = [job_mock]
batch_client.list_namespaced_job.return_value = list_job
batch_client.delete_namespaced_job.return_value = None
monkeypatch.setattr(
"kubernetes.client.BatchV1Api", MagicMock(return_value=batch_client)
)
pod = MagicMock()
pod.metadata.name = "pod_name"
pod.status.phase = "Success"
core_client = MagicMock()
list_pods = MagicMock()
list_pods.items = [pod]
core_client.list_namespaced_pod.return_value = list_pods
monkeypatch.setattr(
"kubernetes.client.CoreV1Api", MagicMock(return_value=core_client)
)
agent = KubernetesAgent()
agent.manage_jobs()
assert batch_client.delete_namespaced_job.called
def test_k8s_agent_generate_deployment_yaml_env_contains_empty_image_pull_secrets(
monkeypatch, cloud_api):
"""
A test to validate that generating the Deployment YAML works correctly if
the IMAGE_PULL_SECRETS env var is an empty string, per issue #5001.
"""
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
monkeypatch.setenv("IMAGE_PULL_SECRETS", "")
agent = KubernetesAgent()
deployment = agent.generate_deployment_yaml(
api="test_api",
namespace="test_namespace",
)
deployment = yaml.safe_load(deployment)
assert "imagePullSecrets" not in deployment["spec"]["template"]["spec"]
agent_env = deployment["spec"]["template"]["spec"]["containers"][0]["env"]
assert agent_env[3]["value"] == ""
def test_k8s_agent_generate_deployment_yaml_labels(monkeypatch, cloud_api):
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
agent = KubernetesAgent()
deployment = agent.generate_deployment_yaml(
token="test_token",
api="test_api",
namespace="test_namespace",
labels=["test_label1", "test_label2"],
)
deployment = yaml.safe_load(deployment)
agent_env = deployment["spec"]["template"]["spec"]["containers"][0]["env"]
assert agent_env[0]["value"] == "test_token"
assert agent_env[1]["value"] == "test_api"
assert agent_env[2]["value"] == "test_namespace"
assert agent_env[4]["value"] == "['test_label1', 'test_label2']"
assert len(deployment["spec"]["template"]["spec"]["containers"]) == 1
def test_k8s_agent_generate_deployment_yaml_local_version(
monkeypatch, version, cloud_api):
monkeypatch.setattr(prefect, "__version__", version[0])
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
agent = KubernetesAgent()
deployment = agent.generate_deployment_yaml(
token="test_token",
api="test_api",
namespace="test_namespace",
resource_manager_enabled=True,
)
deployment = yaml.safe_load(deployment)
agent_yaml = deployment["spec"]["template"]["spec"]["containers"][0]
resource_manager_yaml = deployment["spec"]["template"]["spec"][
"containers"][1]
assert agent_yaml["image"] == "prefecthq/prefect:{}".format(version[1])
assert resource_manager_yaml["image"] == "prefecthq/prefect:{}".format(
version[1])
def test_k8s_agent_includes_agent_labels_in_job(monkeypatch, cloud_api):
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
LocalEnvironment().serialize(),
"id":
"new_id",
"core_version":
"0.13.0",
}),
"id":
"id",
})
agent = KubernetesAgent(labels=["foo", "bar"])
job = agent.generate_job_spec_from_environment(flow_run,
image="test/name:tag")
env = job["spec"]["template"]["spec"]["containers"][0]["env"]
assert env[5]["value"] == "['foo', 'bar']"
def test_k8s_agent_start_max_polls_zero(monkeypatch, runner_token, cloud_api):
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
on_shutdown = MagicMock()
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.on_shutdown",
on_shutdown)
agent_process = MagicMock()
monkeypatch.setattr("prefect.agent.agent.Agent.agent_process",
agent_process)
agent_connect = MagicMock(return_value="id")
monkeypatch.setattr("prefect.agent.agent.Agent.agent_connect",
agent_connect)
heartbeat = MagicMock()
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.heartbeat", heartbeat)
agent = KubernetesAgent(max_polls=0)
agent.start()
assert on_shutdown.call_count == 1
assert agent_process.call_count == 0
assert heartbeat.call_count == 1
def test_k8s_agent_generate_deployment_yaml(monkeypatch, cloud_api):
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
agent = KubernetesAgent()
deployment = agent.generate_deployment_yaml(
key="test-key",
tenant_id="test-tenant",
api="test_api",
namespace="test_namespace",
backend="backend-test",
)
deployment = yaml.safe_load(deployment)
agent_env = deployment["spec"]["template"]["spec"]["containers"][0]["env"]
assert agent_env[0]["value"] == "test-key"
assert agent_env[1]["value"] == "test_api"
assert agent_env[2]["value"] == "test_namespace"
assert agent_env[11]["value"] == "backend-test"
assert agent_env[13] == {
"name": "PREFECT__CLOUD__API_KEY",
"value": "test-key",
}
assert agent_env[14] == {
"name": "PREFECT__CLOUD__TENANT_ID",
"value": "test-tenant",
}
def test_k8s_agent_manage_jobs_client_call(monkeypatch, cloud_api):
gql_return = MagicMock(return_value=MagicMock(data=MagicMock(
set_flow_run_state=None)))
client = MagicMock()
client.return_value.graphql = gql_return
monkeypatch.setattr("prefect.agent.agent.Client", client)
job_mock = MagicMock()
job_mock.metadata.labels = {
"prefect.io/identifier": "id",
"prefect.io/flow_run_id": "fr",
}
job_mock.metadata.name = "my_job"
job_mock.status.failed = False
job_mock.status.succeeded = False
list_job = MagicMock()
list_job.metadata._continue = 0
list_job.items = [job_mock]
pod = MagicMock()
pod.metadata.name = "pod_name"
c_status = MagicMock()
c_status.state.waiting.reason = "ErrImagePull"
pod.status.container_statuses = [c_status]
list_pods = MagicMock()
list_pods.items = [pod]
agent = KubernetesAgent()
agent.batch_client.list_namespaced_job.return_value = list_job
agent.core_client.list_namespaced_pod.return_value = list_pods
agent.manage_jobs()
def test_k8s_agent_manage_jobs_handles_missing_flow_runs(
monkeypatch, cloud_api, caplog):
Client = MagicMock()
Client().get_flow_run_state.side_effect = ObjectNotFoundError()
monkeypatch.setattr("prefect.agent.agent.Client", Client)
job_mock = MagicMock()
job_mock.metadata.labels = {
"prefect.io/identifier": "id",
"prefect.io/flow_run_id": "fr",
}
job_mock.metadata.name = "my_job"
list_job = MagicMock()
list_job.metadata._continue = 0
list_job.items = [job_mock]
pod = MagicMock()
pod.metadata.name = "pod_name"
list_pods = MagicMock()
list_pods.items = [pod]
agent = KubernetesAgent()
agent.batch_client.list_namespaced_job.return_value = list_job
agent.core_client.list_namespaced_pod.return_value = list_pods
agent.heartbeat()
assert (
"Job 'my_job' is for flow run 'fr' which does not exist. It will be ignored."
in caplog.messages)
def test_k8s_agent_removes_yaml_no_volume(monkeypatch, cloud_api):
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
LocalEnvironment().serialize(),
"id":
"id",
"core_version":
"0.13.0",
}),
"id":
"id",
})
agent = KubernetesAgent()
job = agent.generate_job_spec_from_environment(flow_run,
image="test/name:tag")
assert not job["spec"]["template"]["spec"].get("volumes", None)
assert not job["spec"]["template"]["spec"]["containers"][0].get(
"volumeMounts", None)
def test_k8s_agent_generate_deployment_yaml_contains_image_pull_secrets(
monkeypatch, cloud_api
):
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
agent = KubernetesAgent()
deployment = agent.generate_deployment_yaml(
token="test_token",
api="test_api",
namespace="test_namespace",
image_pull_secrets="secrets",
)
deployment = yaml.safe_load(deployment)
assert (
deployment["spec"]["template"]["spec"]["imagePullSecrets"][0]["name"]
== "secrets"
)
agent_env = deployment["spec"]["template"]["spec"]["containers"][0]["env"]
assert agent_env[3]["value"] == "secrets"
def test_k8s_agent_manage_jobs_reports_failed_pods(monkeypatch, cloud_api):
gql_return = MagicMock(
return_value=MagicMock(
data=MagicMock(
set_flow_run_state=None,
write_run_logs=None,
get_flow_run_state=prefect.engine.state.Success(),
)
)
)
client = MagicMock()
client.return_value.graphql = gql_return
monkeypatch.setattr("prefect.agent.agent.Client", client)
job_mock = MagicMock()
job_mock.metadata.labels = {
"prefect.io/identifier": "id",
"prefect.io/flow_run_id": "fr",
}
job_mock.metadata.name = "my_job"
job_mock.status.failed = True
job_mock.status.succeeded = False
batch_client = MagicMock()
list_job = MagicMock()
list_job.metadata._continue = 0
list_job.items = [job_mock]
batch_client.list_namespaced_job.return_value = list_job
monkeypatch.setattr(
"kubernetes.client.BatchV1Api", MagicMock(return_value=batch_client)
)
pod = MagicMock()
pod.metadata.name = "pod_name"
pod.status.phase = "Failed"
terminated = MagicMock()
terminated.exit_code = "code"
terminated.message = "message"
terminated.reason = "reason"
terminated.signal = "signal"
c_status = MagicMock()
c_status.state.terminated = terminated
pod.status.container_statuses = [c_status]
pod2 = MagicMock()
pod2.metadata.name = "pod_name"
pod2.status.phase = "Success"
core_client = MagicMock()
list_pods = MagicMock()
list_pods.items = [pod, pod2]
core_client.list_namespaced_pod.return_value = list_pods
monkeypatch.setattr(
"kubernetes.client.CoreV1Api", MagicMock(return_value=core_client)
)
agent = KubernetesAgent()
agent.manage_jobs()
assert core_client.list_namespaced_pod.called
def test_generate_job_spec_image_pull_secrets_empty_string_in_env(
self, tmpdir, monkeypatch):
"""Regression test for issue #5001."""
run_config = KubernetesRun()
monkeypatch.setenv("IMAGE_PULL_SECRETS", "")
agent = KubernetesAgent(namespace="testing")
job = agent.generate_job_spec(self.build_flow_run(run_config))
assert "imagePullSecrets" not in job["spec"]["template"]["spec"]
def test_generate_job_spec_image_pull_secrets_from_env(
self, tmpdir, monkeypatch):
run_config = KubernetesRun()
monkeypatch.setenv("IMAGE_PULL_SECRETS", "in-env")
agent = KubernetesAgent(namespace="testing")
job = agent.generate_job_spec(self.build_flow_run(run_config))
assert job["spec"]["template"]["spec"]["imagePullSecrets"] == [{
"name":
"in-env"
}]
def test_k8s_agent_manage_pending_pods(monkeypatch, cloud_api):
gql_return = MagicMock(
return_value=MagicMock(
data=MagicMock(set_flow_run_state=None, write_run_logs=None)
)
)
client = MagicMock()
client.return_value.graphql = gql_return
monkeypatch.setattr("prefect.agent.agent.Client", client)
job_mock = MagicMock()
job_mock.metadata.labels = {
"prefect.io/identifier": "id",
"prefect.io/flow_run_id": "fr",
}
job_mock.metadata.name = "my_job"
job_mock.status.failed = False
job_mock.status.succeeded = False
batch_client = MagicMock()
list_job = MagicMock()
list_job.metadata._continue = 0
list_job.items = [job_mock]
batch_client.list_namespaced_job.return_value = list_job
monkeypatch.setattr(
"kubernetes.client.BatchV1Api", MagicMock(return_value=batch_client)
)
dt = pendulum.now()
pod = MagicMock()
pod.metadata.name = "pod_name"
pod.status.phase = "Pending"
event = MagicMock()
event.last_timestamp = dt
event.reason = "reason"
event.message = "message"
core_client = MagicMock()
list_pods = MagicMock()
list_pods.items = [pod]
list_events = MagicMock()
list_events.items = [event]
core_client.list_namespaced_pod.return_value = list_pods
core_client.list_namespaced_event.return_value = list_events
monkeypatch.setattr(
"kubernetes.client.CoreV1Api", MagicMock(return_value=core_client)
)
agent = KubernetesAgent()
agent.manage_jobs()
assert agent.job_pod_event_timestamps["my_job"]["pod_name"] == dt
def test_k8s_agent_generate_deployment_yaml_agent_config_id(
monkeypatch, cloud_api, agent_config_id):
agent = KubernetesAgent()
deployment = yaml.safe_load(
agent.generate_deployment_yaml(agent_config_id=agent_config_id))
cmd_args = deployment["spec"]["template"]["spec"]["containers"][0]["args"]
if agent_config_id:
assert cmd_args == [
"prefect agent kubernetes start --agent-config-id foobar"
]
else:
assert cmd_args == ["prefect agent kubernetes start"]
def test_environment_has_api_key_from_config(self, config_with_api_key):
"""Check that the API key is passed through from the config via environ"""
flow_run = self.build_flow_run(KubernetesRun())
agent = KubernetesAgent(namespace="testing", )
job = agent.generate_job_spec(flow_run)
env_list = job["spec"]["template"]["spec"]["containers"][0]["env"]
env = {item["name"]: item["value"] for item in env_list}
assert env["PREFECT__CLOUD__API_KEY"] == "TEST_KEY"
assert env["PREFECT__CLOUD__AUTH_TOKEN"] == "TEST_KEY"
assert env[
"PREFECT__CLOUD__TENANT_ID"] == config_with_api_key.cloud.tenant_id
def test_k8s_agent_deploy_flow(core_version, command, monkeypatch, cloud_api):
batch_client = MagicMock()
monkeypatch.setattr(
"kubernetes.client.BatchV1Api", MagicMock(return_value=batch_client)
)
core_client = MagicMock()
core_client.list_namespaced_pod.return_value = MagicMock(items=[])
monkeypatch.setattr(
"kubernetes.client.CoreV1Api", MagicMock(return_value=core_client)
)
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
agent = KubernetesAgent()
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": Docker(
registry_url="test", image_name="name", image_tag="tag"
).serialize(),
"environment": LocalEnvironment().serialize(),
"id": "id",
"core_version": core_version,
}
),
"id": "id",
}
)
)
assert agent.batch_client.create_namespaced_job.called
assert (
agent.batch_client.create_namespaced_job.call_args[1]["namespace"] == "default"
)
assert (
agent.batch_client.create_namespaced_job.call_args[1]["body"]["apiVersion"]
== "batch/v1"
)
assert agent.batch_client.create_namespaced_job.call_args[1]["body"]["spec"][
"template"
]["spec"]["containers"][0]["args"] == [command]
def test_environment_has_tenant_id_from_server(self, config_with_api_key):
"""Check that the API key is passed through from the config via environ"""
flow_run = self.build_flow_run(KubernetesRun())
tenant_id = uuid.uuid4()
with set_temporary_config({"cloud.tenant_id": None}):
agent = KubernetesAgent(namespace="testing")
agent.client._get_auth_tenant = MagicMock(return_value=tenant_id)
job = agent.generate_job_spec(flow_run)
env_list = job["spec"]["template"]["spec"]["containers"][0]["env"]
env = {item["name"]: item["value"] for item in env_list}
assert env["PREFECT__CLOUD__API_KEY"] == "TEST_KEY"
assert env["PREFECT__CLOUD__AUTH_TOKEN"] == "TEST_KEY"
assert env["PREFECT__CLOUD__TENANT_ID"] == tenant_id
def test_k8s_agent_generate_deployment_yaml_env_vars(monkeypatch, cloud_api):
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
agent = KubernetesAgent()
env_vars = {"test1": "test2", "test3": "test4"}
deployment = agent.generate_deployment_yaml(env_vars=env_vars)
deployment = yaml.safe_load(deployment)
agent_env = deployment["spec"]["template"]["spec"]["containers"][0]["env"]
assert agent_env[13]["name"] == "PREFECT__CLOUD__AGENT__ENV_VARS"
assert agent_env[13]["value"] == json.dumps(env_vars)
def test_k8s_agent_generate_deployment_yaml_backend_default(monkeypatch, server_api):
c = MagicMock()
monkeypatch.setattr("prefect.agent.agent.Client", c)
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
agent = KubernetesAgent()
deployment = agent.generate_deployment_yaml()
deployment = yaml.safe_load(deployment)
agent_env = deployment["spec"]["template"]["spec"]["containers"][0]["env"]
assert agent_env[11]["value"] == "server"
def test_k8s_agent_deploy_flow_uses_environment_metadata(monkeypatch, cloud_api):
batch_client = MagicMock()
monkeypatch.setattr(
"kubernetes.client.BatchV1Api", MagicMock(return_value=batch_client)
)
core_client = MagicMock()
core_client.list_namespaced_pod.return_value = MagicMock(items=[])
monkeypatch.setattr(
"kubernetes.client.CoreV1Api", MagicMock(return_value=core_client)
)
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
agent = KubernetesAgent()
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": Local().serialize(),
"environment": LocalEnvironment(
metadata={"image": "repo/name:tag"}
).serialize(),
"id": "id",
"core_version": "0.13.0",
}
),
"id": "id",
}
)
)
assert agent.batch_client.create_namespaced_job.called
assert (
agent.batch_client.create_namespaced_job.call_args[1]["body"]["spec"][
"template"
]["spec"]["containers"][0]["image"]
== "repo/name:tag"
)