Exemplo n.º 1
0
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
Exemplo n.º 2
0
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"
Exemplo n.º 3
0
 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"]
Exemplo n.º 4
0
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"
Exemplo n.º 5
0
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()
Exemplo n.º 6
0
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)
Exemplo n.º 7
0
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
Exemplo n.º 8
0
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()
Exemplo n.º 9
0
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
Exemplo n.º 10
0
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"] == ""
Exemplo n.º 11
0
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
Exemplo n.º 12
0
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])
Exemplo n.º 13
0
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']"
Exemplo n.º 14
0
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
Exemplo n.º 15
0
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",
    }
Exemplo n.º 16
0
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()
Exemplo n.º 17
0
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)
Exemplo n.º 18
0
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)
Exemplo n.º 19
0
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"
Exemplo n.º 20
0
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
Exemplo n.º 21
0
 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"]
Exemplo n.º 22
0
 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"
     }]
Exemplo n.º 23
0
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
Exemplo n.º 24
0
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"]
Exemplo n.º 25
0
    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
Exemplo n.º 26
0
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]
Exemplo n.º 27
0
    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
Exemplo n.º 28
0
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)
Exemplo n.º 29
0
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"
Exemplo n.º 30
0
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"
    )