Refactor pause_flow_run for consistency with engine state handling (#…

…7857)

Co-authored-by: Chris Pickett <chris.pickett@prefect.io>
Co-authored-by: Dustin Ngo <dustin.ngo@gmail.com>

src/prefect/exceptions.py

@@ -361,7 +361,3 @@ class NotPausedError(PrefectException):
 
 class FlowPauseTimeout(PrefectException):
     """Raised when a flow pause times out"""
-
-
-class PausedRun(PrefectSignal):
-    """Signal raised when exiting a flow early for nonblocking pauses"""

tests/test_engine.py

@@ -363,23 +363,29 @@ async def test_paused_flows_do_not_block_execution_with_reschedule_flag(
     ):
         frc = partial(FlowRunCreate, deployment_id=deployment.id)
         monkeypatch.setattr("prefect.client.orion.schemas.actions.FlowRunCreate", frc)
+        flow_run_id = None
 
         @task
         async def foo():
             return 42
 
         @flow(task_runner=SequentialTaskRunner())
         async def pausing_flow_without_blocking():
+            nonlocal flow_run_id
+            flow_run_id = get_run_context().flow_run.id
             x = await foo.submit()
             y = await foo.submit()
             await pause_flow_run(timeout=20, reschedule=True)
             z = await foo(wait_for=[x])
             alpha = await foo(wait_for=[y])
             omega = await foo(wait_for=[x, y])
+            assert False, "This line should not be reached"
 
-        flow_run_state = await pausing_flow_without_blocking(return_state=True)
-        assert flow_run_state.is_paused()
-        flow_run_id = flow_run_state.state_details.flow_run_id
+        with pytest.raises(Pause):
+            await pausing_flow_without_blocking(return_state=True)
+
+        flow_run = await orion_client.read_flow_run(flow_run_id)
+        assert flow_run.state.is_paused()
         task_runs = await orion_client.read_task_runs(
             flow_run_filter=FlowRunFilter(id={"any_": [flow_run_id]})
         )
@@ -404,31 +410,36 @@ async def pausing_flow_without_blocking():
             alpha = await foo(wait_for=[y])
             omega = await foo(wait_for=[x, y])
 
-        with pytest.raises(PausedRun):
+        with pytest.raises(Pause):
             await pausing_flow_without_blocking()
 
     async def test_paused_flows_can_be_resumed_then_rescheduled(
         self, orion_client, deployment, monkeypatch
     ):
         frc = partial(FlowRunCreate, deployment_id=deployment.id)
         monkeypatch.setattr("prefect.client.orion.schemas.actions.FlowRunCreate", frc)
+        flow_run_id = None
 
         @task
         async def foo():
             return 42
 
         @flow(task_runner=SequentialTaskRunner())
         async def pausing_flow_without_blocking():
+            nonlocal flow_run_id
+            flow_run_id = get_run_context().flow_run.id
             x = await foo.submit()
             y = await foo.submit()
             await pause_flow_run(timeout=20, reschedule=True)
             z = await foo(wait_for=[x])
             alpha = await foo(wait_for=[y])
             omega = await foo(wait_for=[x, y])
 
-        flow_run_state = await pausing_flow_without_blocking(return_state=True)
-        assert flow_run_state.is_paused()
-        flow_run_id = flow_run_state.state_details.flow_run_id
+        with pytest.raises(Pause):
+            await pausing_flow_without_blocking()
+
+        flow_run = await orion_client.read_flow_run(flow_run_id)
+        assert flow_run.state.is_paused()
 
         await resume_flow_run(flow_run_id)
         flow_run = await orion_client.read_flow_run(flow_run_id)