diff --git a/airflow-core/src/airflow/cli/commands/api_server_command.py b/airflow-core/src/airflow/cli/commands/api_server_command.py index 2e0744be5b8dd..ccc79a1093804 100644 --- a/airflow-core/src/airflow/cli/commands/api_server_command.py +++ b/airflow-core/src/airflow/cli/commands/api_server_command.py @@ -38,7 +38,6 @@ PS = ParamSpec("PS") RT = TypeVar("RT") -AIRFLOW_API_APPS = "AIRFLOW_API_APPS" log = logging.getLogger(__name__) @@ -105,17 +104,17 @@ def with_api_apps_env(func: Callable[[Namespace], RT]) -> Callable[[Namespace], @wraps(func) def wrapper(args: Namespace) -> RT: apps: str = args.apps - original_value = os.environ.get(AIRFLOW_API_APPS) + original_value = os.environ.get("AIRFLOW_API_APPS") try: log.debug("Setting AIRFLOW_API_APPS to: %s", apps) - os.environ[AIRFLOW_API_APPS] = apps + os.environ["AIRFLOW_API_APPS"] = apps return func(args) finally: if original_value is not None: - os.environ[AIRFLOW_API_APPS] = original_value + os.environ["AIRFLOW_API_APPS"] = original_value log.debug("Restored AIRFLOW_API_APPS to: %s", original_value) else: - os.environ.pop(AIRFLOW_API_APPS, None) + os.environ.pop("AIRFLOW_API_APPS", None) log.debug("Removed AIRFLOW_API_APPS from environment") return wrapper diff --git a/airflow-core/tests/unit/cli/commands/test_api_server_command.py b/airflow-core/tests/unit/cli/commands/test_api_server_command.py index 2f3dec30f3dbb..701028c1e8159 100644 --- a/airflow-core/tests/unit/cli/commands/test_api_server_command.py +++ b/airflow-core/tests/unit/cli/commands/test_api_server_command.py @@ -147,15 +147,13 @@ def test_api_apps_env(self, args, dev_mode, original_env): # Verify the AIRFLOW_API_APPS was set correctly if "--apps" in args: - expected_setitem_calls.append( - mock.call(api_server_command.AIRFLOW_API_APPS, parsed_args.apps) - ) + expected_setitem_calls.append(mock.call("AIRFLOW_API_APPS", parsed_args.apps)) # Verify AIRFLOW_API_APPS was cleaned up if original_env is not None: - expected_setitem_calls.append(mock.call(api_server_command.AIRFLOW_API_APPS, original_env)) + expected_setitem_calls.append(mock.call("AIRFLOW_API_APPS", original_env)) else: - mock_environ.pop.assert_called_with(api_server_command.AIRFLOW_API_APPS, None) + mock_environ.pop.assert_called_with("AIRFLOW_API_APPS", None) # Verify that the environment variable was set and cleaned up correctly mock_environ.__setitem__.assert_has_calls(expected_setitem_calls)