def get_flow_run_info(self, flow_run_id, *args, **kwargs):
self.call_count["get_flow_run_info"] += 1
flow_run = self.flow_runs[flow_run_id]
task_runs = [t for t in self.task_runs.values() if t.flow_run_id == flow_run_id]
return FlowRunInfoResult(
id=flow_run.id,
flow_id=flow_run.flow_id,
parameters={},
context=None,
version=flow_run.version,
scheduled_start_time=pendulum.parse("2019-01-25T19:15:58.632412+00:00"),
state=flow_run.state,
task_runs=[
TaskRunInfoResult(
id=tr.id,
task_id=tr.task_slug,
task_slug=tr.task_slug,
version=tr.version,
state=tr.state,
)
for tr in task_runs
],
)
def client(monkeypatch):
cloud_client = MagicMock(
graphql=MagicMock(return_value=MagicMock(data=MagicMock(
flow=[MagicMock(
id="abc123"), MagicMock(id="def456")]))),
create_flow_run=MagicMock(return_value="xyz890"),
get_cloud_url=MagicMock(
return_value="https://api.prefect.io/flow/run/url"),
create_task_run_artifact=MagicMock(return_value="id"),
get_flow_run_info=MagicMock(return_value=FlowRunInfoResult(
id="my-flow-run-id",
name="test-run",
flow_id="xyz890",
version=1,
task_runs=[],
state=state.Success(),
scheduled_start_time=None,
project=ProjectInfo(id="my-project-id", name="Test Project"),
parameters={"test": "ing"},
context={},
)),
)
monkeypatch.setattr("prefect.tasks.prefect.flow_run.Client",
MagicMock(return_value=cloud_client))
monkeypatch.setattr("prefect.Client", MagicMock(return_value=cloud_client))
yield cloud_client
def test_flow_runner_initializes_context_from_cloud(monkeypatch):
from prefect.client.client import FlowRunInfoResult, ProjectInfo
flow = prefect.Flow(name="test")
scheduled_start_time = pendulum.parse("19860920")
get_flow_run_info = MagicMock(return_value=FlowRunInfoResult(
id="my-flow-run-id",
name="my-flow-run-name",
flow_id="my-flow-id",
version=1,
task_runs=[],
state=Pending(),
scheduled_start_time=scheduled_start_time,
project=ProjectInfo(id="my-project-id", name="my-project-name"),
parameters={
"p1": 1,
"p2": 2
},
context={
"c1": 1,
"c2": 2
},
))
client = MagicMock(get_flow_run_info=get_flow_run_info,
set_flow_run_state=MagicMock())
monkeypatch.setattr("prefect.engine.cloud.flow_runner.Client",
MagicMock(return_value=client))
res = CloudFlowRunner(flow=flow).initialize_run(
state=Pending(),
task_states={},
context={
"c2": "two",
"c3": 3
},
task_contexts={},
parameters={
"p2": "two",
"p3": 3
},
)
assert res.context["flow_id"] == "my-flow-id"
assert res.context["flow_run_id"] == "my-flow-run-id"
assert res.context["flow_run_version"] == 1
assert res.context["flow_run_name"] == "my-flow-run-name"
assert res.context["scheduled_start_time"] == scheduled_start_time
assert res.context["project_name"] == "my-project-name"
assert res.context["project_id"] == "my-project-id"
# Explicitly provided parameters override those in cloud
assert res.context["parameters"] == {"p1": 1, "p2": "two", "p3": 3}
# Explicitly provided context overridden by cloud
assert res.context["c1"] == 1
assert res.context["c2"] == 2
assert res.context["c3"] == 3