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