def test_k8s_agent_generate_deployment_yaml_no_resource_manager( 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) 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 len(deployment["spec"]["template"]["spec"]["containers"]) == 1
def test_k8s_agent_generate_deployment_yaml_contains_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", 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_generate_deployment_yaml_env_vars(monkeypatch, runner_token): 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[11]["name"] == "PREFECT__CLOUD__AGENT__ENV_VARS__test1" assert agent_env[11]["value"] == "test2" assert agent_env[12]["name"] == "PREFECT__CLOUD__AGENT__ENV_VARS__test3" assert agent_env[12]["value"] == "test4"
def test_k8s_agent_generate_deployment_yaml_latest(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", 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"
def test_k8s_agent_generate_deployment_yaml_env_vars(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() 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): k8s_config = MagicMock() monkeypatch.setattr("kubernetes.config", k8s_config) 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_generate_deployment_yaml_labels(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", 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_no_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") 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"] == ""
def test_k8s_agent_generate_deployment_yaml_rbac(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", 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"
def test_k8s_agent_generate_deployment_yaml_local_version( monkeypatch, version, runner_token ): monkeypatch.setattr(prefect, "__version__", version[0]) 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_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_generate_deployment_yaml_contains_resources( 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", mem_request="mr", mem_limit="ml", cpu_request="cr", cpu_limit="cl", ) 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"
def test_k8s_agent_generate_deployment_yaml_no_resource_manager(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" ) 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 len(deployment["spec"]["template"]["spec"]["containers"]) == 1