Adjust importlib helper to avoid leaking memory on re-raise (#115377)

This commit is contained in:
J. Nick Koston 2024-04-10 15:19:17 -10:00 committed by GitHub
parent 6bd6adc4f5
commit f0c8c2a684
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 32 additions and 9 deletions

View File

@ -30,11 +30,9 @@ async def async_import_module(hass: HomeAssistant, name: str) -> ModuleType:
if module := cache.get(name):
return module
failure_cache: dict[str, BaseException] = hass.data.setdefault(
DATA_IMPORT_FAILURES, {}
)
if exception := failure_cache.get(name):
raise exception
failure_cache: dict[str, bool] = hass.data.setdefault(DATA_IMPORT_FAILURES, {})
if name in failure_cache:
raise ModuleNotFoundError(f"{name} not found", name=name)
import_futures: dict[str, asyncio.Future[ModuleType]]
import_futures = hass.data.setdefault(DATA_IMPORT_FUTURES, {})
@ -51,7 +49,8 @@ async def async_import_module(hass: HomeAssistant, name: str) -> ModuleType:
module = await hass.async_add_import_executor_job(_get_module, cache, name)
import_future.set_result(module)
except BaseException as ex:
failure_cache[name] = ex
if isinstance(ex, ModuleNotFoundError):
failure_cache[name] = True
import_future.set_exception(ex)
with suppress(BaseException):
# Set the exception retrieved flag on the future since

View File

@ -41,16 +41,40 @@ async def test_async_import_module_failures(hass: HomeAssistant) -> None:
with (
patch(
"homeassistant.helpers.importlib.importlib.import_module",
side_effect=ImportError,
side_effect=ValueError,
),
pytest.raises(ImportError),
pytest.raises(ValueError),
):
await importlib.async_import_module(hass, "test.module")
mock_module = MockModule()
# The failure should be not be cached
with (
patch(
"homeassistant.helpers.importlib.importlib.import_module",
return_value=mock_module,
),
):
assert await importlib.async_import_module(hass, "test.module") is mock_module
async def test_async_import_module_failure_caches_module_not_found(
hass: HomeAssistant,
) -> None:
"""Test importing a module caches ModuleNotFound."""
with (
patch(
"homeassistant.helpers.importlib.importlib.import_module",
side_effect=ModuleNotFoundError,
),
pytest.raises(ModuleNotFoundError),
):
await importlib.async_import_module(hass, "test.module")
mock_module = MockModule()
# The failure should be cached
with (
pytest.raises(ImportError),
pytest.raises(ModuleNotFoundError),
patch(
"homeassistant.helpers.importlib.importlib.import_module",
return_value=mock_module,