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"