def test_all_pipeline_method_steps_are_called(): pipeline = [ "initialize_run", "check_flow_is_pending_or_running", "set_flow_to_running", "get_flow_run_state", ] runner = FlowRunner(Flow(name="test")) for method in pipeline: setattr(runner, method, MagicMock()) # initialize run is unpacked, which MagicMocks dont support runner.initialize_run = MagicMock( return_value=FlowRunnerInitializeResult( MagicMock(), MagicMock(), MagicMock(), MagicMock() ) ) runner.run() for method in pipeline: assert getattr(runner, method).call_count == 1
def test_constant_tasks_arent_submitted_when_mapped(caplog): calls = [] class TrackSubmissions(LocalExecutor): def submit(self, *args, **kwargs): calls.append(kwargs) return super().submit(*args, **kwargs) @prefect.task def add(x): return x + 1 with Flow("constants") as flow: output = add.map([99] * 10) runner = FlowRunner(flow=flow) flow_state = runner.run(return_tasks=[output], executor=TrackSubmissions()) assert flow_state.is_successful() assert flow_state.result[output].result == [100] * 10 ## the add task was submitted 11 times: one for the parent and 10 times for each child assert len(calls) == 11 ## to be safe, ensure '5' isn't in the logs assert len([log.message for log in caplog.records if "99" in log.message]) == 0
def test_constant_tasks_arent_submitted(caplog): calls = [] class TrackSubmissions(LocalExecutor): def submit(self, *args, **kwargs): calls.append(kwargs) return super().submit(*args, **kwargs) @prefect.task def add(x): return x + 1 with Flow("constants") as flow: output = add(5) runner = FlowRunner(flow=flow) flow_state = runner.run(return_tasks=[output], executor=TrackSubmissions()) assert flow_state.is_successful() assert flow_state.result[output].result == 6 ## only add was submitted assert len(calls) == 1 ## to be safe, ensure '5' isn't in the logs assert len([log.message for log in caplog.records if "5" in log.message]) == 0
def test_flow_runner_with_invalid_return_tasks(): flow = Flow(name="test") task = SuccessTask() flow.add_task(task) flow_runner = FlowRunner(flow=flow) state = flow_runner.run(return_tasks=[1]) assert state.is_failed()
def test_flow_runner_runs_basic_flow_with_1_task(): flow = Flow(name="test") task = SuccessTask() flow.add_task(task) flow_runner = FlowRunner(flow=flow) state = flow_runner.run(return_tasks=[task]) assert state == Success(result={task: Success(result=1)})
def test_parameters_overwrite_context(): x = prefect.Parameter("x") f = FlowRunner(Flow(name="test", tasks=[x])) state = f.run( parameters={"x": 2}, context={"parameters": {"x": 5}}, return_tasks=[x] ) assert state.result[x].result == 2
def test_flow_runner_with_no_return_tasks(): """ Make sure FlowRunner accepts return_tasks=None and doesn't raise early error """ flow = Flow(name="test") task = SuccessTask() flow.add_task(task) flow_runner = FlowRunner(flow=flow) assert flow_runner.run(return_tasks=None)
def test_parameters_can_be_set_in_context_if_none_passed(): x = prefect.Parameter("x") f = FlowRunner(Flow(name="test", tasks=[x])) state = f.run(parameters={}, context={"parameters": { "x": 5 }}, return_tasks=[x]) assert state.result[x].result == 5
def test_parameters_overwrite_context_only_if_key_matches(): x = prefect.Parameter("x") y = prefect.Parameter("y") f = FlowRunner(Flow(name="test", tasks=[x, y])) state = f.run( parameters={"x": 2}, context={"parameters": {"x": 5, "y": 6}}, return_tasks=[x, y], )
def test_flow_runner_passes_along_its_run_context_to_tasks(self): @prefect.task def grab_key(): return prefect.context["THE_ANSWER"] with prefect.context(THE_ANSWER=42): runner = FlowRunner(Flow(name="test", tasks=[grab_key])) flow_state = runner.run(return_tasks=[grab_key]) assert flow_state.is_successful() assert flow_state.result[grab_key].result == 42
def test_map_preserves_flowrunners_run_context(executor): @task def whats_id(): return prefect.context.get("special_id") with Flow(name="test-context-preservation") as flow: result = whats_id.map(upstream_tasks=[list(range(10))]) with prefect.context(special_id="FOOBAR"): runner = FlowRunner(flow=flow) flow_state = runner.run(return_tasks=[result]) assert flow_state.is_successful() assert flow_state.result[result].result == ["FOOBAR"] * 10