From 244b190a2df611a45535b9ebb09ddb4452c89bea Mon Sep 17 00:00:00 2001
From: Deepak Kumar <deepakkumar@meta.com>
Date: Sun, 19 Jul 2026 18:14:15 -0700
Subject: [PATCH] edge3: Exit non-zero when worker admin CLI commands fail

The add-worker-queues, remove-worker-queues, and set-worker-concurrency
subcommands raised bare `SystemExit` after catching a TypeError from the
model layer, which yields exit status 0. Shell scripts wrapping these
commands treated failures as success and continued as if the queue or
concurrency change had been applied.

Pass the exception message to SystemExit so the process exits with
status 1 and writes the failure to stderr, matching the file's own
convention for validation-error paths (see the explicit
`raise SystemExit("Error: ...")` calls in the same subcommands) and the
`raise SystemExit(str(e))` pattern already used elsewhere in the edge3
CLI. The model layer already logs the error before raising, so the
redundant logger.error call is removed.
---
 .../providers/edge3/cli/edge_command.py       |  9 ++--
 .../tests/unit/edge3/cli/test_edge_command.py | 53 ++++++++++++++++++-
 2 files changed, 55 insertions(+), 7 deletions(-)

diff --git a/providers/edge3/src/airflow/providers/edge3/cli/edge_command.py b/providers/edge3/src/airflow/providers/edge3/cli/edge_command.py
index 7a9186f2f574e..d1e3133d536ab 100644
--- a/providers/edge3/src/airflow/providers/edge3/cli/edge_command.py
+++ b/providers/edge3/src/airflow/providers/edge3/cli/edge_command.py
@@ -404,8 +404,7 @@ def add_worker_queues(args) -> None:
         add_worker_queues(args.edge_hostname, queues)
         logger.info("Added queues %s to Edge Worker host %s by %s.", queues, args.edge_hostname, getuser())
     except TypeError as e:
-        logger.error(str(e))
-        raise SystemExit
+        raise SystemExit(str(e))
 
 
 @cli_utils.action_cli(check_db=False)
@@ -426,8 +425,7 @@ def remove_worker_queues(args) -> None:
             "Removed queues %s from Edge Worker host %s by %s.", queues, args.edge_hostname, getuser()
         )
     except TypeError as e:
-        logger.error(str(e))
-        raise SystemExit
+        raise SystemExit(str(e))
 
 
 @cli_utils.action_cli(check_db=False)
@@ -450,5 +448,4 @@ def set_remote_worker_concurrency(args) -> None:
             getuser(),
         )
     except TypeError as e:
-        logger.error(str(e))
-        raise SystemExit
+        raise SystemExit(str(e))
diff --git a/providers/edge3/tests/unit/edge3/cli/test_edge_command.py b/providers/edge3/tests/unit/edge3/cli/test_edge_command.py
index 0c0bad84fb9d3..cc9c46dda74f1 100644
--- a/providers/edge3/tests/unit/edge3/cli/test_edge_command.py
+++ b/providers/edge3/tests/unit/edge3/cli/test_edge_command.py
@@ -14,5 +14,56 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from __future__ import annotations
 
-# here is nothing to test, all is CLI wrapper only
+from types import SimpleNamespace
+from unittest import mock
+
+import pytest
+
+from airflow.providers.edge3.cli import edge_command
+
+
+class TestWorkerAdminExitCodes:
+    """CLI worker-admin subcommands must exit non-zero when the model raises."""
+
+    @pytest.fixture(autouse=True)
+    def _stub_db_checks(self):
+        with (
+            mock.patch.object(edge_command, "_check_valid_db_connection"),
+            mock.patch.object(edge_command, "_check_if_registered_edge_host"),
+        ):
+            yield
+
+    @pytest.mark.parametrize(
+        ("cli_func", "model_attr", "args"),
+        [
+            (
+                edge_command.add_worker_queues,
+                "add_worker_queues",
+                SimpleNamespace(edge_hostname="worker-1", queues="q1,q2"),
+            ),
+            (
+                edge_command.remove_worker_queues,
+                "remove_worker_queues",
+                SimpleNamespace(edge_hostname="worker-1", queues="q1"),
+            ),
+            (
+                edge_command.set_remote_worker_concurrency,
+                "set_worker_concurrency",
+                SimpleNamespace(edge_hostname="worker-1", concurrency=8),
+            ),
+        ],
+        ids=["add_worker_queues", "remove_worker_queues", "set_remote_worker_concurrency"],
+    )
+    def test_exits_non_zero_when_model_raises_type_error(self, cli_func, model_attr, args):
+        message = "Cannot mutate worker in OFFLINE state!"
+        with mock.patch(
+            f"airflow.providers.edge3.models.edge_worker.{model_attr}",
+            side_effect=TypeError(message),
+        ):
+            with pytest.raises(SystemExit) as exc_info:
+                cli_func(args)
+        # SystemExit with a string argument yields exit code 1 via sys.exit's
+        # convention; a bare `raise SystemExit` would leave code=None (exit 0).
+        assert exc_info.value.code == message