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
