def test_run_flow(monkeypatch, tmpdir, job_spec_file): environment = KubernetesJobEnvironment(job_spec_file=job_spec_file) flow_runner = MagicMock() flow_runner_class = MagicMock(return_value=flow_runner) monkeypatch.setattr( "prefect.engine.get_default_flow_runner_class", MagicMock(return_value=flow_runner_class), ) d = Local(str(tmpdir)) d.add_flow(prefect.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_class.call_args[1]["flow"].name == "name" assert flow_runner.run.call_args[1]["executor"] is environment.executor
def test_get_flow_returns_flow(): with tempfile.TemporaryDirectory() as tmpdir: storage = Local(directory=tmpdir) f = Flow("test") loc = storage.add_flow(f) runner = storage.get_flow(loc) assert runner == f
def test_environment_execute_calls_callbacks(): start_func = MagicMock() exit_func = MagicMock() with tempfile.TemporaryDirectory() as directory: @prefect.task def add_to_dict(): with open(path.join(directory, "output"), "w") as tmp: tmp.write("success") with open(path.join(directory, "flow_env.prefect"), "w+") as env: flow = prefect.Flow("test", tasks=[add_to_dict]) flow_path = path.join(directory, "flow_env.prefect") with open(flow_path, "wb") as f: cloudpickle.dump(flow, f) environment = RemoteEnvironment(on_start=start_func, on_exit=exit_func) storage = Local(directory) storage.add_flow(flow) environment.execute(flow) with open(path.join(directory, "output"), "r") as file: assert file.read() == "success" assert start_func.called assert exit_func.called
def test_docker_agent_deploy_flow_run_config(api, image_on_run_config): if image_on_run_config: storage = Local() image = "on-run-config" run = DockerRun(image=image, env={"TESTING": "VALUE"}) else: storage = Docker( registry_url="testing", image_name="on-storage", image_tag="tag" ) image = "testing/on-storage:tag" run = DockerRun(env={"TESTING": "VALUE"}) agent = DockerAgent() agent.deploy_flow( flow_run=GraphQLResult( { "flow": GraphQLResult( { "id": "foo", "storage": storage.serialize(), "run_config": run.serialize(), "core_version": "0.13.11", } ), "id": "id", "name": "name", } ) ) assert api.create_container.called assert api.create_container.call_args[0][0] == image assert api.create_container.call_args[1]["environment"]["TESTING"] == "VALUE"
def test_add_flow_raises_if_name_conflict(): with tempfile.TemporaryDirectory() as tmpdir: storage = Local(directory=tmpdir) f = Flow("test") res = storage.add_flow(f) g = Flow("test") with pytest.raises(ValueError, match='name "test"'): storage.add_flow(g)
def test_add_flow_with_weird_name_is_cleaned(): with tempfile.TemporaryDirectory() as tmpdir: s = Local(directory=tmpdir) flow = Flow("WELL what do you know?!~? looks like a test!!!!") loc = s.add_flow(flow) assert "?" not in loc assert "!" not in loc assert " " not in loc assert "~" not in loc
def test_run_flow_no_flow_run_id_in_context(monkeypatch, tmpdir): environment = Environment() d = Local(str(tmpdir)) d.add_flow(Flow("name")) with set_temporary_config({"cloud.auth_token": "test"}): with pytest.raises(ValueError): environment.run_flow()
def test_add_flow_raises_if_name_conflict(): with tempfile.TemporaryDirectory() as tmpdir: storage = Local(directory=tmpdir) f = Flow("test") res = storage.add_flow(f) g = Flow("test") with pytest.raises(ValueError) as exc: storage.add_flow(g) assert 'name "test"' in str(exc.value)
def test_run_flow_no_flow_run_id_in_context(monkeypatch): environment = Environment() with tempfile.TemporaryDirectory() as directory: d = Local(directory) d.add_flow(Flow("name")) with set_temporary_config({"cloud.auth_token": "test"}): with pytest.raises(ValueError): environment.run_flow()
def test_containment(): with tempfile.TemporaryDirectory() as tmpdir: s = Local(directory=tmpdir) f = Flow("test") s.add_flow(f) assert True not in s assert f not in s assert "test" in s assert Flow("other") not in s assert "other" not in s
def build_flow_run(self, config, storage=None): if storage is None: storage = Local() return GraphQLResult({ "flow": GraphQLResult({ "storage": storage.serialize(), "run_config": RunConfigSchema().dump(config), "id": "new_id", "core_version": "0.13.0", }), "id": "id", })
def test_get_flow_from_file_returns_flow(tmpdir): contents = """from prefect import Flow\nf=Flow('test-flow')""" full_path = os.path.join(tmpdir, "flow.py") with open(full_path, "w") as f: f.write(contents) f = Flow("test-flow") storage = Local(stored_as_script=True, path=full_path) storage.add_flow(f) flow = storage.get_flow(full_path) assert flow.run()
def test_add_flow_to_storage(): with tempfile.TemporaryDirectory() as tmpdir: storage = Local(directory=tmpdir) f = Flow("test") assert f.name not in storage res = storage.add_flow(f) assert res.endswith("test.prefect") assert f.name in storage with open(os.path.join(tmpdir, "test.prefect"), "rb") as f: wat = f.read() assert isinstance(wat, bytes) assert cloudpickle.loads(wat).name == "test"
def test_k8s_agent_deploy_flow_raises(monkeypatch, cloud_api): 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() 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
def test_k8s_agent_deploy_flow_uses_environment_metadata( 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_flow(flow_run=GraphQLResult({ "flow": GraphQLResult({ "storage": Local().serialize(), "environment": RemoteEnvironment(metadata={ "image": "repo/name:tag" }).serialize(), "id": "id", }), "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")
def test_docker_agent_deploy_flow_no_pull_using_environment_metadata(api): agent = DockerAgent(no_pull=True) agent.deploy_flow(flow_run=GraphQLResult({ "flow": GraphQLResult({ "id": "foo", "storage": Local().serialize(), "environment": LocalEnvironment(metadata={ "image": "name:tag" }).serialize(), "core_version": "0.13.0", }), "id": "id", "name": "name", })) assert not api.pull.called assert api.create_container.called assert api.start.called
def test_local_agent_deploy_no_existing_python_path(monkeypatch, cloud_api): monkeypatch.delenv("PYTHONPATH", raising=False) 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": "foo", "core_version": "0.13.0", } ), "id": "id", } ) ) assert popen.called assert len(agent.processes) == 1 assert "paths" in popen.call_args[1]["env"]["PYTHONPATH"]
def test_run_flow(monkeypatch): environment = DaskKubernetesEnvironment() flow_runner = MagicMock() monkeypatch.setattr( "prefect.engine.get_default_flow_runner_class", MagicMock(return_value=flow_runner), ) kube_cluster = MagicMock() monkeypatch.setattr("dask_kubernetes.KubeCluster", kube_cluster) with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "flow_env.prefect"), "w+") as env: storage = Local(directory) flow = prefect.Flow("test", storage=storage) flow_path = os.path.join(directory, "flow_env.prefect") with open(flow_path, "wb") as f: cloudpickle.dump(flow, f) with set_temporary_config({"cloud.auth_token": "test"}): with prefect.context(flow_file_path=os.path.join( directory, "flow_env.prefect")): environment.run_flow() assert flow_runner.call_args[1]["flow"].name == "test"
async def flow(self, project_id, tmpdir): """ A simple diamond flow """ @prefect.task def numbers(): return [1, 2, 3] @prefect.task def add(x, y): return x + y with prefect.Flow( "simple map", environment=LocalEnvironment(), storage=Local(directory=tmpdir), ) as flow: flow.numbers1 = numbers() flow.numbers2 = numbers() flow.add = add.map(flow.numbers1, flow.numbers2) add_state_handler(flow) with set_temporary_config(key="dev", value=True): flow.server_id = await api.flows.create_flow( project_id=project_id, serialized_flow=flow.serialize(build=True)) return flow
async def flow(self, project_id, tmpdir): """ A simple diamond flow """ @prefect.task def numbers(): return [1, 2, 3] @prefect.task def add(x): if x == 2: raise prefect.engine.signals.FAIL("Don't run for 2") elif not isinstance(x, int): return -99 return x + 10 with prefect.Flow( "simple map", environment=LocalEnvironment(), storage=Local(directory=tmpdir), ) as flow: flow.numbers = numbers() flow.add1 = add.map(flow.numbers) flow.add2 = add.map(flow.add1) flow.add2.trigger = prefect.triggers.all_finished add_state_handler(flow) with set_temporary_config(key="dev", value=True): flow.server_id = await api.flows.create_flow( project_id=project_id, serialized_flow=flow.serialize(build=True)) return flow
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
def test_docker_agent_deploy_flow_unsupported_run_config(api): agent = DockerAgent() with pytest.raises( TypeError, match= "`run_config` of type `LocalRun`, only `DockerRun` is supported", ): 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_docker_agent_deploy_flow_no_pull_using_environment_metadata( monkeypatch, cloud_api): 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(no_pull=True) agent.deploy_flow(flow_run=GraphQLResult({ "flow": GraphQLResult({ "id": "foo", "storage": Local().serialize(), "environment": LocalEnvironment(metadata={ "image": "name:tag" }).serialize(), "core_version": "0.13.0", }), "id": "id", "name": "name", })) assert not api.pull.called assert api.create_container.called assert api.start.called
def test_docker_agent_deploy_flow_storage_raises(monkeypatch, cloud_api): monkeypatch.setattr("prefect.agent.agent.Client", MagicMock()) 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() with pytest.raises(ValueError): agent.deploy_flow(flow_run=GraphQLResult({ "flow": GraphQLResult({ "storage": Local().serialize(), "id": "foo", "environment": LocalEnvironment().serialize(), "core_version": "0.13.0", }), "id": "id", "name": "name", "version": "version", })) assert not api.pull.called
def test_deploy_flows_require_docker_storage(monkeypatch, runner_token): boto3_client = MagicMock() boto3_client.describe_task_definition.return_value = {"tags": []} boto3_client.run_task.return_value = {} monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client)) with pytest.raises(Exception) as excinfo: agent = FargateAgent() agent.deploy_flow( GraphQLResult({ "flow": GraphQLResult({ "storage": Local().serialize(), "id": "id", "version": 2, "name": "name", }), "id": "id", "name": "name", })) assert boto3_client.describe_task_definition.not_called assert boto3_client.run_task.not_called
def test_run_flow_calls_callbacks(monkeypatch): start_func = MagicMock() exit_func = MagicMock() file_path = os.path.dirname( prefect.environments.execution.dask.k8s.__file__) environment = KubernetesJobEnvironment(path.join(file_path, "job.yaml"), on_start=start_func, on_exit=exit_func) flow_runner = MagicMock() monkeypatch.setattr( "prefect.engine.get_default_flow_runner_class", MagicMock(return_value=flow_runner), ) with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "flow_env.prefect"), "w+") as env: storage = Local(directory) flow = prefect.Flow("test", storage=storage) flow_path = os.path.join(directory, "flow_env.prefect") with open(flow_path, "wb") as f: cloudpickle.dump(flow, f) with set_temporary_config({"cloud.auth_token": "test"}): with prefect.context(flow_file_path=os.path.join( directory, "flow_env.prefect")): environment.run_flow() assert flow_runner.call_args[1]["flow"].name == "test" assert start_func.called assert exit_func.called
def test_docker_agent_deploy_flow_storage_raises(monkeypatch, api): monkeypatch.setattr("prefect.agent.agent.Client", MagicMock()) agent = DockerAgent() with pytest.raises(ValueError): agent.deploy_flow( flow_run=GraphQLResult( { "flow": GraphQLResult( { "storage": Local().serialize(), "id": "foo", "environment": LocalEnvironment().serialize(), "core_version": "0.13.0", } ), "id": "id", "name": "name", "version": "version", } ) ) assert not api.pull.called
def test_create_local_storage_without_validation(): storage = Local(directory="C:\\Users\\chris\\.prefect\\flows", validate=False) assert storage assert storage.directory == "C:\\Users\\chris\\.prefect\\flows" assert isinstance(storage.result, LocalResult) assert storage.result.dir == storage.directory
def test_docker_agent_deploy_flow_uses_environment_metadata(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"
def test_run_flow_calls_callbacks(monkeypatch): start_func = MagicMock() exit_func = MagicMock() environment = DaskKubernetesEnvironment(on_start=start_func, on_exit=exit_func) flow_runner = MagicMock() monkeypatch.setattr( "prefect.engine.get_default_flow_runner_class", MagicMock(return_value=flow_runner), ) kube_cluster = MagicMock() monkeypatch.setattr("dask_kubernetes.KubeCluster", kube_cluster) with tempfile.TemporaryDirectory() as directory: d = Local(directory) d.add_flow(prefect.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.dask.k8s.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 start_func.called assert exit_func.called