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