Esempio n. 1
0
def test_k8s_agent_deploy_flow(core_version, command, monkeypatch, cloud_api):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    batch_client = MagicMock()
    batch_client.create_namespaced_job.return_value = {}
    monkeypatch.setattr("kubernetes.client.BatchV1Api",
                        MagicMock(retrurn_value=batch_client))

    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]
Esempio n. 2
0
def test_k8s_agent_start_max_polls_zero(monkeypatch, runner_token, cloud_api):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    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 == 0
Esempio n. 3
0
def test_k8s_agent_removes_yaml_no_volume(monkeypatch, cloud_api):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    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.replace_job_spec_yaml(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)
Esempio n. 4
0
def test_k8s_agent_includes_agent_labels_in_job(monkeypatch, cloud_api):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    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.replace_job_spec_yaml(flow_run, image="test/name:tag")
    env = job["spec"]["template"]["spec"]["containers"][0]["env"]
    assert env[5]["value"] == "['foo', 'bar']"
Esempio n. 5
0
def test_k8s_agent_replace_yaml_responds_to_logging_config(
    monkeypatch, cloud_api, flag
):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    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.replace_job_spec_yaml(flow_run, image="test/name:tag")
    env = job["spec"]["template"]["spec"]["containers"][0]["env"]
    assert env[6]["value"] == str(not flag).lower()
Esempio n. 6
0
def test_k8s_agent_generate_deployment_yaml(monkeypatch, runner_token):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    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_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 resource_manager_env[0]["value"] == "test_token"
    assert resource_manager_env[1]["value"] == "test_api"
    assert resource_manager_env[3]["value"] == "test_namespace"
Esempio n. 7
0
def test_k8s_agent_start_max_polls(monkeypatch, runner_token, cloud_api):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    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=1)
    agent.start()

    assert agent_process.called
    assert heartbeat.called
Esempio n. 8
0
def test_k8s_agent_manage_jobs_pass(monkeypatch, cloud_api):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    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()
Esempio n. 9
0
def test_k8s_agent_generate_deployment_yaml_labels(monkeypatch, cloud_api):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    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
Esempio n. 10
0
def test_k8s_agent_generate_deployment_yaml_contains_image_pull_secrets(
        monkeypatch, cloud_api):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    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"
Esempio n. 11
0
def test_k8s_agent_generate_deployment_yaml_latest(monkeypatch, cloud_api):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    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,
        latest=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:latest"
    assert resource_manager_yaml["image"] == "prefecthq/prefect:latest"
Esempio n. 12
0
def test_k8s_agent_replace_yaml_no_pull_secrets(monkeypatch, cloud_api):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    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.replace_job_spec_yaml(flow_run,
                                      image="test/name:tag",
                                      identifier="identifier")

    assert not job["spec"]["template"]["spec"].get("imagePullSecrets", None)
Esempio n. 13
0
def test_k8s_agent_replace_yaml_responds_to_logging_config(
        monkeypatch, runner_token, flag):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    flow_run = GraphQLResult({
        "flow":
        GraphQLResult({
            "storage":
            Docker(registry_url="test", image_name="name",
                   image_tag="tag").serialize(),
            "id":
            "new_id",
        }),
        "id":
        "id",
        "name":
        "name",
    })

    with set_temporary_config({
            "cloud.agent.auth_token": "token",
            "logging.log_to_cloud": flag
    }):
        agent = KubernetesAgent()
        job = agent.replace_job_spec_yaml(flow_run)
        env = job["spec"]["template"]["spec"]["containers"][0]["env"]
        assert env[6]["value"] == str(flag).lower()
Esempio n. 14
0
def test_k8s_agent_generate_deployment_yaml_contains_resources(monkeypatch, cloud_api):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    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"
Esempio n. 15
0
def test_k8s_agent_manage_jobs_delete_jobs(monkeypatch, cloud_api):

    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

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

    agent = KubernetesAgent()
    agent.manage_jobs()

    assert batch_client.delete_namespaced_job.called
Esempio n. 16
0
def test_k8s_agent_deploy_flow_raises(monkeypatch, cloud_api):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    batch_client = MagicMock()
    batch_client.create_namespaced_job.return_value = {}
    monkeypatch.setattr(
        "kubernetes.client.BatchV1Api", MagicMock(retrurn_value=batch_client)
    )

    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
Esempio n. 17
0
def test_k8s_agent_deploy_flow_uses_environment_metadata(
        monkeypatch, runner_token):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    batch_client = MagicMock()
    batch_client.create_namespaced_job.return_value = {}
    monkeypatch.setattr("kubernetes.client.BatchV1Api",
                        MagicMock(retrurn_value=batch_client))

    agent = KubernetesAgent()
    agent.deploy_flow(flow_run=GraphQLResult({
        "flow":
        GraphQLResult({
            "storage":
            Local().serialize(),
            "environment":
            RemoteEnvironment(metadata={
                "image": "repo/name:tag"
            }).serialize(),
            "id":
            "id",
        }),
        "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")
Esempio n. 18
0
def test_k8s_agent_deploy_flows(monkeypatch, runner_token):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    batch_client = MagicMock()
    batch_client.create_namespaced_job.return_value = {}
    monkeypatch.setattr("kubernetes.client.BatchV1Api",
                        MagicMock(retrurn_value=batch_client))

    agent = KubernetesAgent()
    agent.deploy_flows(flow_runs=[
        GraphQLResult({
            "flow":
            GraphQLResult({
                "storage":
                Docker(registry_url="test", image_name="name",
                       image_tag="tag").serialize(),
                "id":
                "id",
            }),
            "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")
Esempio n. 19
0
def test_k8s_agent_replace_yaml(monkeypatch, runner_token, cloud_api):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    monkeypatch.setenv("IMAGE_PULL_SECRETS", "my-secret")
    monkeypatch.setenv("JOB_MEM_REQUEST", "mr")
    monkeypatch.setenv("JOB_MEM_LIMIT", "ml")
    monkeypatch.setenv("JOB_CPU_REQUEST", "cr")
    monkeypatch.setenv("JOB_CPU_LIMIT", "cl")

    flow_run = GraphQLResult({
        "flow":
        GraphQLResult({
            "storage":
            Docker(registry_url="test", image_name="name",
                   image_tag="tag").serialize(),
            "environment":
            RemoteEnvironment().serialize(),
            "id":
            "new_id",
        }),
        "id":
        "id",
    })

    with set_temporary_config({
            "cloud.agent.auth_token": "token",
            "logging.log_to_cloud": True
    }):
        agent = KubernetesAgent()
        job = agent.replace_job_spec_yaml(flow_run, image="test/name:tag")

        assert job["metadata"]["labels"]["prefect.io/flow_run_id"] == "id"
        assert job["metadata"]["labels"]["prefect.io/flow_id"] == "new_id"
        assert (job["spec"]["template"]["metadata"]["labels"]
                ["prefect.io/flow_run_id"] == "id")
        assert (job["spec"]["template"]["spec"]["containers"][0]["image"] ==
                "test/name:tag")

        env = job["spec"]["template"]["spec"]["containers"][0]["env"]

        assert env[0]["value"] == "https://api.prefect.io"
        assert env[1]["value"] == "token"
        assert env[2]["value"] == "id"
        assert env[3]["value"] == "new_id"
        assert env[4]["value"] == "default"
        assert env[5]["value"] == "[]"
        assert env[6]["value"] == "true"

        assert (job["spec"]["template"]["spec"]["imagePullSecrets"][0]["name"]
                == "my-secret")

        resources = job["spec"]["template"]["spec"]["containers"][0][
            "resources"]
        assert resources["requests"]["memory"] == "mr"
        assert resources["limits"]["memory"] == "ml"
        assert resources["requests"]["cpu"] == "cr"
        assert resources["limits"]["cpu"] == "cl"
Esempio n. 20
0
def test_k8s_agent_replace_yaml_uses_user_env_vars(monkeypatch, runner_token,
                                                   cloud_api):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    monkeypatch.setenv("IMAGE_PULL_SECRETS", "my-secret")
    monkeypatch.setenv("JOB_MEM_REQUEST", "mr")
    monkeypatch.setenv("JOB_MEM_LIMIT", "ml")
    monkeypatch.setenv("JOB_CPU_REQUEST", "cr")
    monkeypatch.setenv("JOB_CPU_LIMIT", "cl")

    flow_run = GraphQLResult({
        "flow":
        GraphQLResult({
            "storage":
            Docker(registry_url="test", image_name="name",
                   image_tag="tag").serialize(),
            "environment":
            RemoteEnvironment().serialize(),
            "id":
            "new_id",
        }),
        "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.replace_job_spec_yaml(flow_run, image="test/name:tag")

        assert job["metadata"]["labels"]["prefect.io/flow_run_id"] == "id"
        assert job["metadata"]["labels"]["prefect.io/flow_id"] == "new_id"
        assert (job["spec"]["template"]["metadata"]["labels"]
                ["prefect.io/flow_run_id"] == "id")
        assert (job["spec"]["template"]["spec"]["containers"][0]["image"] ==
                "test/name:tag")

        env = job["spec"]["template"]["spec"]["containers"][0]["env"]

        assert env[0]["value"] == "https://api.prefect.io"
        assert env[1]["value"] == "token"
        assert env[2]["value"] == "id"
        assert env[3]["value"] == "new_id"
        assert env[4]["value"] == "default"
        assert env[5]["value"] == "[]"
        assert env[6]["value"] == "true"

        user_vars = [
            dict(name="AUTH_THING", value="foo"),
            dict(name="PKG_SETTING", value="bar"),
        ]
        assert env[-1] in user_vars
        assert env[-2] in user_vars
Esempio n. 21
0
def test_k8s_agent_generate_deployment_yaml_backend_default(monkeypatch, server_api):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    agent = KubernetesAgent()
    deployment = agent.generate_deployment_yaml()

    deployment = yaml.safe_load(deployment)

    agent_env = deployment["spec"]["template"]["spec"]["containers"][0]["env"]

    assert agent_env[9]["value"] == "server"
Esempio n. 22
0
def test_k8s_agent_generate_deployment_yaml_no_image_pull_secrets(
        monkeypatch, runner_token):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    agent = KubernetesAgent()
    deployment = agent.generate_deployment_yaml(token="test_token",
                                                api="test_api",
                                                namespace="test_namespace")

    deployment = yaml.safe_load(deployment)

    assert deployment["spec"]["template"]["spec"].get(
        "imagePullSecrets") is None
Esempio n. 23
0
def test_k8s_agent_heartbeat_creates_file(monkeypatch, runner_token):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    with tempfile.TemporaryDirectory() as tempdir:

        monkeypatch.setattr(
            "prefect.agent.kubernetes.agent.AGENT_DIRECTORY",
            "{}/.prefect/agent".format(tempdir),
        )

        agent = KubernetesAgent()
        agent.heartbeat()

        assert path.exists("{}/.prefect/agent/heartbeat".format(tempdir))
Esempio n. 24
0
def test_k8s_agent_init(monkeypatch, runner_token):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    agent = KubernetesAgent()
    assert agent
    assert agent.batch_client
Esempio n. 25
0
def test_k8s_agent_generate_deployment_yaml_rbac(monkeypatch, cloud_api):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    agent = KubernetesAgent()
    deployment = agent.generate_deployment_yaml(
        token="test_token", api="test_api", namespace="test_namespace", rbac=True
    )

    deployment = yaml.safe_load_all(deployment)

    for document in deployment:
        if "rbac" in document:
            assert "rbac" in document["apiVersion"]
            assert document["metadata"]["namespace"] == "test_namespace"
            assert document["metadata"]["name"] == "prefect-agent-rbac"
Esempio n. 26
0
def test_k8s_agent_generate_deployment_yaml_no_image_pull_secrets(
    monkeypatch, cloud_api
):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    agent = KubernetesAgent()
    deployment = agent.generate_deployment_yaml(
        token="test_token", api="test_api", namespace="test_namespace"
    )

    deployment = yaml.safe_load(deployment)

    assert deployment["spec"]["template"]["spec"].get("imagePullSecrets") is None
    agent_env = deployment["spec"]["template"]["spec"]["containers"][0]["env"]
    assert agent_env[3]["value"] == ""
Esempio n. 27
0
def test_k8s_agent_deploy_flow_uses_environment_metadata(monkeypatch, cloud_api):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    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"
    )
Esempio n. 28
0
def install(label, env, **kwargs):
    """Generate a supervisord.conf file for a Local agent"""
    from prefect.agent.kubernetes import KubernetesAgent

    deployment = KubernetesAgent.generate_deployment_yaml(
        labels=sorted(set(label)), env_vars=dict(e.split("=", 2) for e in env), **kwargs
    )
    click.echo(deployment)
Esempio n. 29
0
def test_k8s_agent_generate_deployment_yaml_env_vars(monkeypatch, cloud_api):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    agent = KubernetesAgent()
    deployment = agent.generate_deployment_yaml(
        env_vars={"test1": "test2", "test3": "test4"}
    )

    deployment = yaml.safe_load(deployment)

    agent_env = deployment["spec"]["template"]["spec"]["containers"][0]["env"]

    assert agent_env[13]["name"] == "PREFECT__CLOUD__AGENT__ENV_VARS__test1"
    assert agent_env[13]["value"] == "test2"
    assert agent_env[14]["name"] == "PREFECT__CLOUD__AGENT__ENV_VARS__test3"
    assert agent_env[14]["value"] == "test4"
Esempio n. 30
0
def test_k8s_agent_init(monkeypatch, cloud_api):
    k8s_config = MagicMock()
    monkeypatch.setattr("kubernetes.config", k8s_config)

    agent = KubernetesAgent()
    assert agent
    assert agent.labels == []
    assert agent.name == "agent"
    assert agent.batch_client