Repository: PrefectHQ/prefect

Issue:
Alter signature of Client.set_flow_run_state
## Current behavior
*Please describe how the feature works today*
Currently, the `Client.set_flow_run_state(flow_run_id, state)` call always returns `None`. Unlike the task equivalent, which returns the new `State` object. All places where the client is calling this method are implicitly handling the state change by creating the new state, calling the method, then proceeding with that new state. 



## Proposed behavior
*Please describe your proposed change to the current behavior*
The `set_flow_run_state` call should return the new `State` for a couple reasons. First, and most importantly, in order to extend the `Flow`'s options of states (instead of Pending -> Running -> Success / Failed), we need to be able to implement logic of listening to Cloud's response to the state change event. Secondly and less important, returning the new `State` for both the task and flow client calls would be nice for consistency reasons.



## Example
*Please give an example of how the enhancement would be useful*
While this feature itself won't change the current behavior, it will allow extension of Flow state control in the future!


Discussion:

Resolving patch:
diff --git a/src/prefect/client/client.py b/src/prefect/client/client.py
--- a/src/prefect/client/client.py
+++ b/src/prefect/client/client.py
@@ -2,6 +2,7 @@
 import json
 import os
 import re
+import time
 import uuid
 import warnings
 from pathlib import Path
@@ -10,7 +11,6 @@
 
 import pendulum
 import toml
-import time
 from slugify import slugify
 
 import prefect
@@ -952,7 +952,7 @@ def update_task_run_heartbeat(self, task_run_id: str) -> None:
 
     def set_flow_run_state(
         self, flow_run_id: str, version: int, state: "prefect.engine.state.State"
-    ) -> None:
+    ) -> "prefect.engine.state.State":
         """
         Sets new state for a flow run in the database.
 
@@ -961,6 +961,9 @@ def set_flow_run_state(
             - version (int): the current version of the flow run state
             - state (State): the new state for this flow run
 
+        Returns:
+            - State: the state the current flow run should be considered in
+
         Raises:
             - ClientError: if the GraphQL mutation is bad for any reason
         """
@@ -987,6 +990,8 @@ def set_flow_run_state(
             ),
         )  # type: Any
 
+        return state
+
     def get_latest_cached_states(
         self, task_id: str, cache_key: Optional[str], created_after: datetime.datetime
     ) -> List["prefect.engine.state.State"]:
diff --git a/src/prefect/engine/cloud/flow_runner.py b/src/prefect/engine/cloud/flow_runner.py
--- a/src/prefect/engine/cloud/flow_runner.py
+++ b/src/prefect/engine/cloud/flow_runner.py
@@ -114,7 +114,7 @@ def call_runner_target_handlers(self, old_state: State, new_state: State) -> Sta
 
         try:
             cloud_state = new_state
-            self.client.set_flow_run_state(
+            state = self.client.set_flow_run_state(
                 flow_run_id=flow_run_id, version=version, state=cloud_state
             )
         except Exception as exc:


Test patch:
diff --git a/tests/client/test_client.py b/tests/client/test_client.py
--- a/tests/client/test_client.py
+++ b/tests/client/test_client.py
@@ -11,7 +11,7 @@
 import prefect
 from prefect.client.client import Client, FlowRunInfoResult, TaskRunInfoResult
 from prefect.engine.result import NoResult, Result, SafeResult
-from prefect.engine.state import Pending
+from prefect.engine.state import Pending, State
 from prefect.utilities.configuration import set_temporary_config
 from prefect.utilities.exceptions import AuthorizationError, ClientError
 from prefect.utilities.graphql import GraphQLResult, decompress
@@ -670,10 +670,11 @@ def test_set_flow_run_state(patch_post):
         {"cloud.api": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
     ):
         client = Client()
-    result = client.set_flow_run_state(
-        flow_run_id="74-salt", version=0, state=Pending()
-    )
-    assert result is None
+
+    state = Pending()
+    result = client.set_flow_run_state(flow_run_id="74-salt", version=0, state=state)
+    assert isinstance(result, State)
+    assert isinstance(result, Pending)
 
 
 def test_set_flow_run_state_with_error(patch_post):
diff --git a/tests/engine/cloud/test_cloud_flow_runner.py b/tests/engine/cloud/test_cloud_flow_runner.py
--- a/tests/engine/cloud/test_cloud_flow_runner.py
+++ b/tests/engine/cloud/test_cloud_flow_runner.py
@@ -9,15 +9,15 @@
 import prefect
 from prefect.client.client import Client, FlowRunInfoResult
 from prefect.engine.cloud import CloudFlowRunner, CloudTaskRunner
-from prefect.engine.signals import LOOP
 from prefect.engine.result import NoResult, Result, SafeResult
-from prefect.engine.results import PrefectResult, SecretResult
 from prefect.engine.result_handlers import (
     ConstantResultHandler,
     JSONResultHandler,
     ResultHandler,
     SecretResultHandler,
 )
+from prefect.engine.results import PrefectResult, SecretResult
+from prefect.engine.signals import LOOP
 from prefect.engine.state import (
     Cancelled,
     Failed,
@@ -52,7 +52,9 @@ def cloud_settings():
 def client(monkeypatch):
     cloud_client = MagicMock(
         get_flow_run_info=MagicMock(return_value=MagicMock(state=None, parameters={})),
-        set_flow_run_state=MagicMock(),
+        set_flow_run_state=MagicMock(
+            side_effect=lambda flow_run_id, version, state: state
+        ),
         get_task_run_info=MagicMock(return_value=MagicMock(state=None)),
         set_task_run_state=MagicMock(
             side_effect=lambda task_run_id, version, state, cache_for: state
@@ -540,7 +542,6 @@ def downstream(l):
     client.get_flow_run_info = MagicMock(
         return_value=MagicMock(context={"task_loop_count": 20})
     )
-    client.set_flow_run_state = MagicMock()
 
     flow_state = CloudFlowRunner(flow=f).run(return_tasks=[inter, final])
 
@@ -604,6 +605,9 @@ def write_run_logs(self, *args, **kwargs):
         def set_task_run_state(self, *args, **kwargs):
             return kwargs.get("state")
 
+        def set_flow_run_state(self, *args, **kwargs):
+            return kwargs.get("state")
+
         def get_flow_run_info(self, *args, **kwargs):
             return MagicMock(
                 id="flow_run_id",
diff --git a/tests/engine/cloud/test_cloud_flows.py b/tests/engine/cloud/test_cloud_flows.py
--- a/tests/engine/cloud/test_cloud_flows.py
+++ b/tests/engine/cloud/test_cloud_flows.py
@@ -171,6 +171,7 @@ def set_flow_run_state(self, flow_run_id, version, state, **kwargs):
             fr.version += 1
         else:
             raise ValueError("Invalid flow run update")
+        return state
 
     def set_task_run_state(self, task_run_id, version, state, **kwargs):
         self.call_count["set_task_run_state"] += 1