def test_flow_run_task_with_no_matching_flow(self, client, server_api): # verify a ValueError is raised if the client returns no flows task = StartFlowRun(flow_name="Test Flow", project_name="Demo") client.graphql = MagicMock(return_value=MagicMock(data=MagicMock( flow=[]))) with pytest.raises(ValueError, match="Flow 'Test Flow' not found."): task.run()
def test_flow_run_task_with_wait(self, client, server_api): # verify that create_flow_run was called task = StartFlowRun( flow_name="Test Flow", project_name="Demo", parameters={"test": "ing"}, run_name="test-run", wait=True, poll_interval=timedelta(seconds=3), ) assert task.poll_interval == timedelta(seconds=3) # Run flow, and assert that signals a success with pytest.raises(signals.SUCCESS) as exc_info: task.run() flow_state_signal = exc_info.value assert isinstance(flow_state_signal.state, state.Success) # Check flow ID assert str(flow_state_signal).split(" ")[0] == "xyz890" # verify the GraphQL query was called with the correct arguments query_args = list( client.graphql.call_args_list[0][0][0]["query"].keys())[0] assert "Test Flow" in query_args # verify create_flow_run was called with the correct arguments client.create_flow_run.assert_called_once_with( flow_id="abc123", parameters={"test": "ing"}, idempotency_key=None, context=None, run_name="test-run", scheduled_start_time=None, run_config=None, )
def test_flow_run_link_artifact(self, client, cloud_api): task = StartFlowRun( project_name="Test Project", flow_name="Test Flow", parameters={"test": "ing"}, run_name="test-run", ) with prefect.context(running_with_backend=True, task_run_id="trid"): task.run() client.create_task_run_artifact.assert_called_once_with( data={"link": "/flow/run/url"}, kind="link", task_run_id="trid" )
def test_flow_run_task(self, client, cloud_api): # verify that create_flow_run was called task = StartFlowRun( project_name="Test Project", flow_name="Test Flow", parameters={"test": "ing"}, run_name="test-run", ) # verify that run returns the new flow run ID assert task.run() == "xyz890" # verify the GraphQL query was called with the correct arguments query_args = list( client.graphql.call_args_list[0][0][0]["query"].keys())[0] assert "Test Project" in query_args assert "Test Flow" in query_args # verify create_flow_run was called with the correct arguments client.create_flow_run.assert_called_once_with( flow_id="abc123", parameters={"test": "ing"}, idempotency_key=None, context=None, run_name="test-run", scheduled_start_time=None, )
def test_flow_run_task_submit_args(self, client, cloud_api, idempotency_key, task_run_id): run_config = UniversalRun() # verify that create_flow_run was called task = StartFlowRun( project_name="Test Project", flow_name="Test Flow", parameters={"test": "ing"}, run_config=run_config, run_name="test-run", ) # verify that run returns the new flow run ID with prefect.context(task_run_id=task_run_id): assert task.run(idempotency_key=idempotency_key) == "xyz890" # verify the GraphQL query was called with the correct arguments query_args = list( client.graphql.call_args_list[0][0][0]["query"].keys())[0] assert "Test Project" in query_args assert "Test Flow" in query_args # verify create_flow_run was called with the correct arguments assert client.create_flow_run.call_args[1] == dict( flow_id="abc123", parameters={"test": "ing"}, run_config=run_config, idempotency_key=idempotency_key or task_run_id, context=None, run_name="test-run", scheduled_start_time=None, )
def test_idempotency_key_uses_map_index_if_present(self, client, cloud_api): # verify that create_flow_run was called task = StartFlowRun(project_name="Test Project", flow_name="Test Flow") # verify that run returns the new flow run ID with prefect.context(flow_run_id="test-id", map_index=4): assert task.run() == "xyz890" # verify create_flow_run was called with the correct arguments client.create_flow_run.assert_called_once_with( flow_id="abc123", idempotency_key="test-id-4", parameters=None, context=None, run_name=None, )
def test_flow_run_task_uses_scheduled_start_time(self, client, cloud_api): in_one_hour = pendulum.now().add(hours=1) # verify that create_flow_run was called task = StartFlowRun( project_name="Test Project", flow_name="Test Flow", scheduled_start_time=in_one_hour, ) # verify that run returns the new flow run ID assert task.run() == "xyz890" # verify create_flow_run was called with the correct arguments client.create_flow_run.assert_called_once_with( flow_id="abc123", parameters=None, idempotency_key=None, context=None, run_name=None, scheduled_start_time=in_one_hour, )
def test_flow_run_task_with_flow_run_id(self, client, cloud_api): # verify that create_flow_run was called task = StartFlowRun( project_name="Test Project", flow_name="Test Flow", parameters={"test": "ing"}, ) # verify that run returns the new flow run ID with prefect.context(flow_run_id="test-id"): assert task.run() == "xyz890" # verify create_flow_run was called with the correct arguments client.create_flow_run.assert_called_once_with( flow_id="abc123", parameters={"test": "ing"}, idempotency_key="test-id", context=None, run_name=None, )
def create_prefect_flow_run(flow_name: str, project_name: str, task_refs: List, params: Mapping) -> str: """Creates new prefect flow run for given flow id, parameters, task references and API server URL to send GraphQL requests to. Returns results value and state from a Prefect flow run. """ try: flow_run = StartFlowRun(flow_name=flow_name, project_name=project_name, parameters=params) flow_run_id = flow_run.run() client = Client() while True: time.sleep(10) flow_run_info = client.get_flow_run_info(flow_run_id) flow_state = flow_run_info.state task_runs_info = flow_run_info.task_runs if flow_state.is_finished(): task_res_locs = {} for task_run in task_runs_info: # Return ref if ref string is a substring of any task slug ref = next((ref_str for ref_str in task_refs if ref_str in task_run.task_slug), None) if ref: task_id = task_run.id task_state = client.get_task_run_state(task_id) task_res_locs[ref] = task_state._result.location task_results = {} for ref, loc in task_res_locs.items(): local_res = LocalResult() result = local_res.read(loc) task_results[ref] = result.value return task_results, flow_state, task_res_locs except ValueError as err: raise err
def test_flow_run_task_without_flow_name(self, server_api): # verify that a ValueError is raised without a flow name task = StartFlowRun() with pytest.raises(ValueError, match="Must provide a flow name."): task.run()
def test_flow_run_task_without_project_name(self, cloud_api): # verify that a ValueError is raised without a project name task = StartFlowRun(flow_name="Test Flow") with pytest.raises(ValueError, match="Must provide a project name."): task.run()