From 194cc9341bc311b546d2d44363430684c3ce0e9a Mon Sep 17 00:00:00 2001
From: Chong Shen Ng <chong.shen@flower.ai>
Date: Mon, 16 Dec 2024 20:34:10 +0000
Subject: [PATCH] refactor(framework) Add `try`-`except` block to `flwr stop`
 for `JSON` output (#4712)

---
 src/py/flwr/cli/stop.py | 52 ++++++++++++++++++++++++-----------------
 1 file changed, 30 insertions(+), 22 deletions(-)

diff --git a/src/py/flwr/cli/stop.py b/src/py/flwr/cli/stop.py
index acfa818efb32..84002f0f6e27 100644
--- a/src/py/flwr/cli/stop.py
+++ b/src/py/flwr/cli/stop.py
@@ -48,34 +48,42 @@ def stop(
     ] = None,
 ) -> None:
     """Stop a run."""
-    # Load and validate federation config
-    typer.secho("Loading project configuration... ", fg=typer.colors.BLUE)
+    try:
+        # Load and validate federation config
+        typer.secho("Loading project configuration... ", fg=typer.colors.BLUE)
 
-    pyproject_path = app / FAB_CONFIG_FILE if app else None
-    config, errors, warnings = load_and_validate(path=pyproject_path)
-    config = process_loaded_project_config(config, errors, warnings)
-    federation, federation_config = validate_federation_in_project_config(
-        federation, config
-    )
-    exit_if_no_address(federation_config, "stop")
+        pyproject_path = app / FAB_CONFIG_FILE if app else None
+        config, errors, warnings = load_and_validate(path=pyproject_path)
+        config = process_loaded_project_config(config, errors, warnings)
+        federation, federation_config = validate_federation_in_project_config(
+            federation, config
+        )
+        exit_if_no_address(federation_config, "stop")
 
-    try:
-        auth_plugin = try_obtain_cli_auth_plugin(app, federation)
-        channel = init_channel(app, federation_config, auth_plugin)
-        stub = ExecStub(channel)  # pylint: disable=unused-variable # noqa: F841
+        try:
+            auth_plugin = try_obtain_cli_auth_plugin(app, federation)
+            channel = init_channel(app, federation_config, auth_plugin)
+            stub = ExecStub(channel)  # pylint: disable=unused-variable # noqa: F841
 
-        typer.secho(f"✋ Stopping run ID {run_id}...", fg=typer.colors.GREEN)
-        _stop_run(stub, run_id=run_id)
+            typer.secho(f"✋ Stopping run ID {run_id}...", fg=typer.colors.GREEN)
+            _stop_run(stub, run_id=run_id)
 
-    except ValueError as err:
-        typer.secho(
-            f"❌ {err}",
-            fg=typer.colors.RED,
-            bold=True,
-        )
+        except ValueError as err:
+            typer.secho(
+                f"❌ {err}",
+                fg=typer.colors.RED,
+                bold=True,
+            )
+            raise typer.Exit(code=1) from err
+        finally:
+            channel.close()
+    except (
+        typer.Exit,
+        Exception,
+    ) as err:  # pylint: disable=broad-except, W0612  # noqa: F841
         raise typer.Exit(code=1) from err
     finally:
-        channel.close()
+        pass
 
 
 def _stop_run(