def test_run_flow(monkeypatch):
    environment = CloudEnvironment()

    flow_runner = MagicMock()
    monkeypatch.setattr(
        "prefect.engine.get_default_flow_runner_class",
        MagicMock(return_value=flow_runner),
    )
    monkeypatch.setattr(
        "prefect.environments.execution.cloud.environment.sys.exit",
        MagicMock())

    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:
            flow = prefect.Flow("test")
            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"
示例#2
0
def test_run_flow(monkeypatch):
    environment = CloudEnvironment()

    flow_runner = MagicMock()
    monkeypatch.setattr("prefect.engine.FlowRunner", 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:
            flow = prefect.Flow("test")
            flow_path = os.path.join(directory, "flow_env.prefect")
            with open(flow_path, "w") as f:
                json.dump(flow.serialize(), 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"