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" )