From ceb615f68bd59d1203b9d8da0671b6ad9507c0c5 Mon Sep 17 00:00:00 2001 From: CooperWang0912 Date: Thu, 23 Jul 2026 18:04:07 +0800 Subject: [PATCH] Improved Wrapper Logic --- archinstall/lib/global_menu.py | 29 ++++--- tests/test_global_menu_ui.py | 154 --------------------------------- 2 files changed, 16 insertions(+), 167 deletions(-) delete mode 100644 tests/test_global_menu_ui.py diff --git a/archinstall/lib/global_menu.py b/archinstall/lib/global_menu.py index b21e4c14..35a5eb40 100644 --- a/archinstall/lib/global_menu.py +++ b/archinstall/lib/global_menu.py @@ -1,5 +1,6 @@ import inspect from collections.abc import Callable +from functools import wraps from typing import Any, override from archinstall.default_profiles.profile import GreeterType @@ -101,22 +102,24 @@ class GlobalMenu(AbstractMenu[None]): prefix = self._get_status_prefix(item) item.text = f'{prefix}{base_title}' - def _wrap_action(self, key: str, action: Callable[..., Any]) -> Callable[..., Any]: - async def wrapper(*args, **kwargs) -> Any: - if inspect.iscoroutinefunction(action): - result = await action(*args, **kwargs) - else: - result = action(*args, **kwargs) - if inspect.isawaitable(result): - result = await result + def _wrap_action(self, item: MenuItem, action: Callable[..., Any]) -> Callable[..., Any]: + @wraps(action) + async def wrapper(*args: Any, **kwargs: Any) -> Any: + try: + if inspect.iscoroutinefunction(action): + result = await action(*args, **kwargs) + else: + result = action(*args, **kwargs) + if inspect.isawaitable(result): + result = await result - item = self._item_group.find_by_key(key) - if item: + # Directly update the item's value item.value = result - self._update_item_labels() + return result - return result + finally: + self._update_item_labels() return wrapper @@ -254,7 +257,7 @@ class GlobalMenu(AbstractMenu[None]): if wrap_actions: for item in menu_options: if item.key and item.action: - item.action = self._wrap_action(item.key, item.action) + item.action = self._wrap_action(item, item.action) return menu_options diff --git a/tests/test_global_menu_ui.py b/tests/test_global_menu_ui.py deleted file mode 100644 index cb1435d3..00000000 --- a/tests/test_global_menu_ui.py +++ /dev/null @@ -1,154 +0,0 @@ -import asyncio -import inspect -from unittest.mock import MagicMock - - -class MockSpecialMenuKey: - SAVE = 'save' - INSTALL = 'install' - ABORT = 'abort' - - -class MockMenuItem: - def __init__(self, key=None, read_only=False, mandatory=False, value=None): - self.key = key - self.read_only = read_only - self.mandatory = mandatory - self.value = value - - def has_value(self): - return self.value is not None - - -def isolated_get_status_prefix(item: MockMenuItem) -> str: - special_keys = (MockSpecialMenuKey.SAVE, MockSpecialMenuKey.INSTALL, MockSpecialMenuKey.ABORT) - - if item.read_only or item.key in special_keys: - return '' - - if item.key == 'auth_config': - auth_config = item.value - has_root = getattr(auth_config, 'root_enc_password', None) is not None - has_super = getattr(auth_config, 'has_superuser', lambda: False)() if auth_config else False - - if auth_config is not None and (has_root or has_super): - return '[bold green][✓][/bold green] ' - return '[bold red][!][/bold red] ' - - if item.has_value(): - return '[bold green][✓][/bold green] ' - elif item.mandatory: - return '[bold red][!][/bold red] ' - else: - return '[bold yellow][•][/bold yellow] ' - - -def isolated_wrap_action(item_dictionary, update_callback, key, action): - - async def wrapper(*args, **kwargs): - if inspect.iscoroutinefunction(action): - result = await action(*args, **kwargs) - else: - result = action(*args, **kwargs) - if inspect.isawaitable(result): - result = await result - - if key in item_dictionary: - item_dictionary[key].value = result - - update_callback() - return result - - return wrapper - - -class TestPrefixLogic: - def test_special_keys_have_no_prefix(self): - item = MockMenuItem(key=MockSpecialMenuKey.SAVE) - assert isolated_get_status_prefix(item) == '' - - def test_read_only_has_no_prefix(self): - item = MockMenuItem(read_only=True) - assert isolated_get_status_prefix(item) == '' - - def test_configured_item_has_checkmark(self): - item = MockMenuItem(key='hostname', value='archlinux') - prefix = isolated_get_status_prefix(item) - assert '[✓]' in prefix - assert 'green' in prefix - - def test_missing_mandatory_item_has_warning(self): - item = MockMenuItem(key='disk_config', mandatory=True, value=None) - prefix = isolated_get_status_prefix(item) - assert '[!]' in prefix - assert 'red' in prefix - - def test_missing_optional_item_has_dot(self): - item = MockMenuItem(key='network_config', mandatory=False, value=None) - prefix = isolated_get_status_prefix(item) - assert '[•]' in prefix - assert 'yellow' in prefix - - -class TestAuthConfigLogic: - def test_auth_missing_entirely(self): - item = MockMenuItem(key='auth_config', value=None) - assert '[!]' in isolated_get_status_prefix(item) - - def test_auth_invalid_state(self): - mock_auth = MagicMock() - mock_auth.root_enc_password = None - mock_auth.has_superuser.return_value = False - - item = MockMenuItem(key='auth_config', value=mock_auth) - assert '[!]' in isolated_get_status_prefix(item) - - def test_auth_valid_root(self): - mock_auth = MagicMock() - mock_auth.root_enc_password = 'hashed' - mock_auth.has_superuser.return_value = False - - item = MockMenuItem(key='auth_config', value=mock_auth) - assert '[✓]' in isolated_get_status_prefix(item) - - def test_auth_valid_superuser(self): - mock_auth = MagicMock() - mock_auth.root_enc_password = None - mock_auth.has_superuser.return_value = True - - item = MockMenuItem(key='auth_config', value=mock_auth) - assert '[✓]' in isolated_get_status_prefix(item) - - -class TestActionWrapperLogic: - def test_async_action_updates_value(self): - mock_item = MockMenuItem(key='disk') - item_dict = {'disk': mock_item} - update_spy = MagicMock() - - async def dummy_async_action(): - return 'new_disk_layout' - - wrapped = isolated_wrap_action(item_dict, update_spy, 'disk', dummy_async_action) - - result = asyncio.run(wrapped()) - - assert result == 'new_disk_layout' - assert mock_item.value == 'new_disk_layout' - update_spy.assert_called_once() - - def test_sync_action_updates_value(self): - mock_item = MockMenuItem(key='hostname') - item_dict = {'hostname': mock_item} - update_spy = MagicMock() - - def dummy_sync_action(): - return 'my-pc' - - wrapped = isolated_wrap_action(item_dict, update_spy, 'hostname', dummy_sync_action) - - result = asyncio.run(wrapped()) - - assert result == 'my-pc' - assert mock_item.value == 'my-pc' - update_spy.assert_called_once()