def test_populate_env_vars_uses_user_provided_env_vars(runner_token):
with set_temporary_config(
{
"cloud.api": "api",
"logging.log_to_cloud": True,
"cloud.agent.auth_token": "token",
}
):
agent = LocalAgent(env_vars=dict(AUTH_THING="foo"))
env_vars = agent.populate_env_vars(GraphQLResult({"id": "id"}))
assert env_vars["AUTH_THING"] == "foo"
def test_populate_env_vars_from_agent_config():
agent = LocalAgent(env_vars=dict(AUTH_THING="foo"))
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}))
assert env_vars["AUTH_THING"] == "foo"
def test_populate_env_vars_removes_none_values():
agent = LocalAgent(env_vars=dict(MISSING_VAR=None))
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}))
assert "MISSING_VAR" not in env_vars
def test_docker_agent_networks_as_modes_can_be_overriden_by_run_config(
api, network):
api.create_networking_config.return_value = {network: "config1"}
api.create_endpoint_config.return_value = "endpoint-config"
agent = DockerAgent(networks=[network])
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"id":
"foo",
"name":
"flow-name",
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"core_version":
"0.13.0",
}),
"run_config":
DockerRun(host_config={
"network_mode": "foobar"
}).serialize(),
"id":
"id",
"name":
"name",
}))
assert network in agent.networks
api.create_networking_config.assert_called_once_with(
{network: "endpoint-config"})
_, container_create_kwargs = api.create_container.call_args
assert container_create_kwargs["networking_config"] == {network: "config1"}
_, host_config_kwargs = api.create_host_config.call_args
assert host_config_kwargs["network_mode"] == "foobar"
def test_docker_agent_deploy_flow_show_flow_logs(api, monkeypatch):
process = MagicMock()
monkeypatch.setattr("multiprocessing.Process", process)
agent = DockerAgent(show_flow_logs=True)
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"id":
"foo",
"name":
"flow-name",
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"core_version":
"0.13.0",
}),
"id":
"id",
"name":
"name",
}))
process_kwargs = dict(
target=_stream_container_logs,
kwargs={
"base_url": agent.base_url,
"container_id": "container_id",
"timeout": 60,
},
)
process.assert_called_with(**process_kwargs)
# Check all arguments to `multiprocessing.Process` are pickleable
assert pickle.loads(pickle.dumps(process_kwargs)) == process_kwargs
assert len(agent.processes) == 1
assert api.create_container.called
assert api.start.called
def deploy_flow(self, run_config, **kwargs):
agent = ECSAgent(**kwargs)
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Local().serialize(),
"run_config":
run_config.serialize() if run_config else None,
"id":
"flow-id",
"version":
1,
"name":
"Test Flow",
"core_version":
"0.13.0",
}),
"id":
"flow-run-id",
})
return agent.deploy_flow(flow_run)
def test_docker_agent_deploy_flow_run_config(api, run_kind,
has_docker_storage):
if has_docker_storage:
storage = Docker(registry_url="testing",
image_name="on-storage",
image_tag="tag")
image = "testing/on-storage:tag"
else:
storage = Local()
image = ("on-run-config" if run_kind == "docker" else
"prefecthq/prefect:all_extras-0.13.11")
if run_kind == "docker":
env = {"TESTING": "VALUE"}
run = DockerRun(image=image, env=env)
else:
env = {}
run = None if run_kind == "missing" else UniversalRun()
agent = DockerAgent()
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"id": "foo",
"storage": storage.serialize(),
"run_config": run.serialize() if run else None,
"core_version": "0.13.11",
}),
"id":
"id",
"name":
"name",
}))
assert api.create_container.called
assert api.create_container.call_args[0][0] == image
res_env = api.create_container.call_args[1]["environment"]
for k, v in env.items():
assert res_env[k] == v
def test_k8s_agent_includes_agent_labels_in_job(monkeypatch, runner_token):
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": "id",
}
),
"id": "id",
}
)
agent = KubernetesAgent(labels=["foo", "bar"])
job = agent.replace_job_spec_yaml(flow_run)
env = job["spec"]["template"]["spec"]["containers"][0]["env"]
assert env[4]["value"] == "['foo', 'bar']"
def test_populate_env_vars_sets_log_to_cloud(flag):
agent = LocalAgent(no_cloud_logs=flag)
assert agent.log_to_cloud is not flag
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}))
assert env_vars["PREFECT__CLOUD__SEND_FLOW_RUN_LOGS"] == str(
not flag).lower()
def test_nomad_agent_deploy_flows(monkeypatch, runner_token):
requests = MagicMock()
monkeypatch.setattr("prefect.agent.nomad.agent.requests", requests)
agent = NomadAgent()
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 requests.post.called
assert requests.post.call_args[1]["json"]
def test_k8s_agent_replace_yaml_no_pull_secrets(monkeypatch, runner_token):
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": "id",
}
),
"id": "id",
}
)
agent = KubernetesAgent()
job = agent.replace_job_spec_yaml(flow_run)
assert not job["spec"]["template"]["spec"]["imagePullSecrets"][0]["name"]
def test_local_agent_responds_to_logging_config(runner_token, flag):
with set_temporary_config({
"cloud.agent.auth_token": "TEST_TOKEN",
"logging.log_to_cloud": flag
}):
agent = LocalAgent()
assert agent.log_to_cloud is flag
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name"
}))
assert env_vars["PREFECT__LOGGING__LOG_TO_CLOUD"] == str(flag).lower()
def test_local_agent_deploy_processes_valid_storage(storage, monkeypatch):
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
agent = LocalAgent()
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": storage.serialize(),
"id": "foo",
"core_version": "0.13.0",
}
),
"id": "id",
}
)
)
assert popen.called
assert len(agent.processes) == 1
def test_run_flow(monkeypatch):
environment = Environment()
flow_runner = MagicMock()
monkeypatch.setattr(
"prefect.engine.get_default_flow_runner_class",
MagicMock(return_value=flow_runner),
)
executor = MagicMock()
monkeypatch.setattr(
"prefect.engine.get_default_executor_class",
MagicMock(return_value=executor),
)
with tempfile.TemporaryDirectory() as directory:
d = Local(directory)
d.add_flow(Flow("name"))
gql_return = MagicMock(return_value=MagicMock(data=MagicMock(flow_run=[
GraphQLResult({
"flow":
GraphQLResult({
"name": "name",
"storage": d.serialize(),
})
})
], )))
client = MagicMock()
client.return_value.graphql = gql_return
monkeypatch.setattr("prefect.environments.execution.base.Client",
client)
with set_temporary_config({"cloud.auth_token": "test"
}), prefect.context({"flow_run_id": "id"}):
environment.run_flow()
assert flow_runner.call_args[1]["flow"].name == "name"
assert executor.call_args == None
def test_populate_env_vars_no_existing_python_path(monkeypatch):
monkeypatch.delenv("PYTHONPATH", raising=False)
agent = LocalAgent(import_paths=["paths"])
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}))
assert "paths" in env_vars["PYTHONPATH"]
def test_local_agent_deploy_processes_gitlab_storage(monkeypatch, cloud_api):
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
agent = LocalAgent()
gitlab = GitLab("test/repo", path="path/to/flow.py")
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": gitlab.serialize(),
"id": "foo",
}
),
"id": "id",
}
)
)
assert popen.called
assert len(agent.processes) == 1
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
def test_populate_env_vars_includes_agent_labels():
agent = LocalAgent(labels=["42", "marvin"])
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}))
expected = str(["42", "marvin"] + DEFAULT_AGENT_LABELS)
assert env_vars["PREFECT__CLOUD__AGENT__LABELS"] == expected
def test_local_agent_deploy_keep_existing_python_path(monkeypatch, runner_token):
monkeypatch.setenv("PYTHONPATH", "cool:python:path")
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
agent = LocalAgent(import_paths=["paths"])
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult({"storage": Local(directory="test").serialize()}),
"id": "id",
}
)
)
python_path = popen.call_args[1]["env"]["PYTHONPATH"]
assert popen.called
assert len(agent.processes) == 1
assert "cool:python:path" in python_path
assert "paths" in python_path
def test_k8s_agent_replace_yaml(monkeypatch, runner_token):
k8s_config = MagicMock()
monkeypatch.setattr("kubernetes.config", k8s_config)
monkeypatch.setenv("IMAGE_PULL_SECRETS", "my-secret")
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"id":
"id",
}),
"id":
"id",
})
with set_temporary_config({"cloud.agent.auth_token": "token"}):
agent = KubernetesAgent()
job = agent.replace_job_spec_yaml(flow_run)
assert job["metadata"]["labels"]["flow_run_id"] == "id"
assert job["metadata"]["labels"]["flow_id"] == "id"
assert job["spec"]["template"]["metadata"]["labels"][
"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"] == "default"
assert (job["spec"]["template"]["spec"]["imagePullSecrets"][0]["name"]
== "my-secret")
def test_populate_env_vars_from_agent_config(api):
agent = DockerAgent(env_vars=dict(AUTH_THING="foo"))
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}), "test-image")
assert env_vars["AUTH_THING"] == "foo"
def test_k8s_agent_deploy_flows_continues(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": Local().serialize(),
"id": "id"
}),
"id": "id",
})
])
assert not agent.batch_client.create_namespaced_job.called
def test_docker_agent_deploy_flow_unsupported_run_config(api):
agent = DockerAgent()
with pytest.raises(TypeError,
match="Unsupported RunConfig type: LocalRun"):
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"storage": Local().serialize(),
"run_config": LocalRun().serialize(),
"id": "foo",
"core_version": "0.13.0",
}),
"id":
"id",
"name":
"name",
"version":
"version",
}))
assert not api.pull.called
def test_deploy_flow_raises(monkeypatch, runner_token):
boto3_client = MagicMock()
boto3_client.describe_task_definition.return_value = {}
boto3_client.run_task.return_value = {"tasks": [{"taskArn": "test"}]}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
agent = FargateAgent()
with pytest.raises(ValueError):
agent.deploy_flow(flow_run=GraphQLResult(
{
"flow": GraphQLResult({
"storage": Local().serialize(),
"id": "id"
}),
"id": "id",
}))
assert not boto3_client.describe_task_definition.called
assert not boto3_client.run_task.called
def test_get_key_value_calls_client_query_correctly(
self, monkeypatch, cloud_api):
key_value_id_gql_return = GraphQLResult(
data=dict(key_value=[GraphQLResult({"id": "123"})], ))
delete_key_value_gql_return = GraphQLResult(
data=dict(delete_key_value=GraphQLResult({"success": True}), ))
# helper function to return key value id
# and the delete_key_value depending on input
def fake_graphql_responses(*args, **kwargs):
if "query" in kwargs["query"]:
return key_value_id_gql_return
return delete_key_value_gql_return
Client = MagicMock()
Client.return_value.graphql.side_effect = fake_graphql_responses
monkeypatch.setattr("prefect.backend.kv_store.Client", Client)
success = delete_key(key="foo")
assert success
assert Client.return_value.graphql.call_args_list == [
call(query={
"query": {
'key_value(where: { key: { _eq: "foo" } })': {"id"}
}
}),
call(
query={
"mutation($input: delete_key_value_input!)": {
"delete_key_value(input: $input)": {"success"}
}
},
variables={"input": {
"key_value_id": "123"
}},
),
]
def test_environment_has_agent_token_from_config(api, config_with_token):
agent = DockerAgent()
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}), "test-image")
assert env_vars["PREFECT__CLOUD__AUTH_TOKEN"] == "TEST_TOKEN"
def test_populate_env_vars(api, backend):
agent = DockerAgent()
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}), "test-image")
if backend == "server":
cloud_api = "http://host.docker.internal:4200"
else:
cloud_api = prefect.config.cloud.api
expected_vars = {
"PREFECT__BACKEND":
backend,
"PREFECT__CLOUD__API":
cloud_api,
"PREFECT__CLOUD__AUTH_TOKEN":
"",
"PREFECT__CLOUD__API_KEY":
"",
"PREFECT__CLOUD__TENANT_ID":
"",
"PREFECT__CLOUD__AGENT__LABELS":
"[]",
"PREFECT__CONTEXT__FLOW_RUN_ID":
"id",
"PREFECT__CONTEXT__FLOW_ID":
"foo",
"PREFECT__CONTEXT__IMAGE":
"test-image",
"PREFECT__CLOUD__USE_LOCAL_SECRETS":
"false",
"PREFECT__CLOUD__SEND_FLOW_RUN_LOGS":
"true",
"PREFECT__LOGGING__LOG_TO_CLOUD":
"true",
"PREFECT__LOGGING__LEVEL":
"INFO",
"PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS":
"prefect.engine.cloud.CloudFlowRunner",
"PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS":
"prefect.engine.cloud.CloudTaskRunner",
}
assert env_vars == expected_vars
def test_populate_env_vars_includes_agent_labels(cloud_api):
with set_temporary_config({
"cloud.api": "api",
"logging.log_to_cloud": True,
"cloud.agent.auth_token": "token",
}):
agent = LocalAgent(labels=["42", "marvin"])
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}))
expected_vars = {
"PREFECT__CLOUD__API":
"api",
"PREFECT__CLOUD__AUTH_TOKEN":
"token",
"PREFECT__CLOUD__AGENT__LABELS":
str([
"42",
"marvin",
socket.gethostname(),
"azure-flow-storage",
"gcs-flow-storage",
"s3-flow-storage",
"github-flow-storage",
"webhook-flow-storage",
"gitlab-flow-storage",
]),
"PREFECT__CONTEXT__FLOW_RUN_ID":
"id",
"PREFECT__CONTEXT__FLOW_ID":
"foo",
"PREFECT__CLOUD__USE_LOCAL_SECRETS":
"false",
"PREFECT__LOGGING__LOG_TO_CLOUD":
"true",
"PREFECT__LOGGING__LEVEL":
"INFO",
"PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS":
"prefect.engine.cloud.CloudFlowRunner",
"PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS":
"prefect.engine.cloud.CloudTaskRunner",
}
assert env_vars == expected_vars
def test_docker_agent_deploy_flow_show_flow_logs(monkeypatch, cloud_api):
process = MagicMock()
monkeypatch.setattr("multiprocessing.Process", process)
api = MagicMock()
api.ping.return_value = True
api.create_container.return_value = {"Id": "container_id"}
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
agent = DockerAgent(show_flow_logs=True)
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"id":
"foo",
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
LocalEnvironment().serialize(),
"core_version":
"0.13.0",
}),
"id":
"id",
"name":
"name",
}))
process.assert_called_with(target=agent.stream_container_logs,
kwargs={"container_id": "container_id"})
assert len(agent.processes) == 1
assert api.create_container.called
assert api.start.called
def test_docker_agent_deploy_flow_uses_environment_metadata(monkeypatch, cloud_api):
api = MagicMock()
api.ping.return_value = True
api.create_container.return_value = {"Id": "container_id"}
api.create_host_config.return_value = {"AutoRemove": True}
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
agent = DockerAgent()
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"id": "foo",
"storage": Local().serialize(),
"environment": LocalEnvironment(
metadata={"image": "repo/name:tag"}
).serialize(),
"core_version": "0.13.0",
}
),
"id": "id",
"name": "name",
}
)
)
assert api.pull.called
assert api.create_container.called
assert api.start.called
assert api.create_host_config.call_args[1]["auto_remove"] is True
assert api.create_container.call_args[1]["command"] == "prefect execute flow-run"
assert api.create_container.call_args[1]["host_config"]["AutoRemove"] is True
assert api.start.call_args[1]["container"] == "container_id"
