Ensure asyncio blocking checks are undone after tests run (#119542)

* Ensure asyncio blocking checks are undone after tests run

* no reason to ever enable twice

* we are patching objects, make it more generic

* make sure bootstrap unblocks as well

* move disable to tests only

* re-protect

* Update tests/test_block_async_io.py

Co-authored-by: Erik Montnemery <erik@montnemery.com>

* Revert "Update tests/test_block_async_io.py"

This reverts commit 2d46028e21.

* tweak name

* fixture only

* Update tests/conftest.py

* Update tests/conftest.py

* Apply suggestions from code review

---------

Co-authored-by: Erik Montnemery <erik@montnemery.com>

homeassistant/block_async_io.py

@@ -1,7 +1,9 @@
 """Block blocking calls being done in asyncio."""
 
 import builtins
+from collections.abc import Callable
 from contextlib import suppress
+from dataclasses import dataclass
 import glob
 from http.client import HTTPConnection
 import importlib
@@ -46,53 +48,131 @@ def _check_sleep_call_allowed(mapped_args: dict[str, Any]) -> bool:
     return False
 
 
+@dataclass(slots=True, frozen=True)
+class BlockingCall:
+    """Class to hold information about a blocking call."""
+
+    original_func: Callable
+    object: object
+    function: str
+    check_allowed: Callable[[dict[str, Any]], bool] | None
+    strict: bool
+    strict_core: bool
+    skip_for_tests: bool
+
+
+_BLOCKING_CALLS: tuple[BlockingCall, ...] = (
+    BlockingCall(
+        original_func=HTTPConnection.putrequest,
+        object=HTTPConnection,
+        function="putrequest",
+        check_allowed=None,
+        strict=True,
+        strict_core=True,
+        skip_for_tests=False,
+    ),
+    BlockingCall(
+        original_func=time.sleep,
+        object=time,
+        function="sleep",
+        check_allowed=_check_sleep_call_allowed,
+        strict=True,
+        strict_core=True,
+        skip_for_tests=False,
+    ),
+    BlockingCall(
+        original_func=glob.glob,
+        object=glob,
+        function="glob",
+        check_allowed=None,
+        strict=False,
+        strict_core=False,
+        skip_for_tests=False,
+    ),
+    BlockingCall(
+        original_func=glob.iglob,
+        object=glob,
+        function="iglob",
+        check_allowed=None,
+        strict=False,
+        strict_core=False,
+        skip_for_tests=False,
+    ),
+    BlockingCall(
+        original_func=os.walk,
+        object=os,
+        function="walk",
+        check_allowed=None,
+        strict=False,
+        strict_core=False,
+        skip_for_tests=False,
+    ),
+    BlockingCall(
+        original_func=os.listdir,
+        object=os,
+        function="listdir",
+        check_allowed=None,
+        strict=False,
+        strict_core=False,
+        skip_for_tests=True,
+    ),
+    BlockingCall(
+        original_func=os.scandir,
+        object=os,
+        function="scandir",
+        check_allowed=None,
+        strict=False,
+        strict_core=False,
+        skip_for_tests=True,
+    ),
+    BlockingCall(
+        original_func=builtins.open,
+        object=builtins,
+        function="open",
+        check_allowed=_check_file_allowed,
+        strict=False,
+        strict_core=False,
+        skip_for_tests=True,
+    ),
+    BlockingCall(
+        original_func=importlib.import_module,
+        object=importlib,
+        function="import_module",
+        check_allowed=_check_import_call_allowed,
+        strict=False,
+        strict_core=False,
+        skip_for_tests=True,
+    ),
+)
+
+
+@dataclass(slots=True)
+class BlockedCalls:
+    """Class to track which calls are blocked."""
+
+    calls: set[BlockingCall]
+
+
+_BLOCKED_CALLS = BlockedCalls(set())
+
+
 def enable() -> None:
     """Enable the detection of blocking calls in the event loop."""
+    calls = _BLOCKED_CALLS.calls
+    if calls:
+        raise RuntimeError("Blocking call detection is already enabled")
+
     loop_thread_id = threading.get_ident()
-    # Prevent urllib3 and requests doing I/O in event loop
-    HTTPConnection.putrequest = protect_loop(  # type: ignore[method-assign]
-        HTTPConnection.putrequest, loop_thread_id=loop_thread_id
-    )
+    for blocking_call in _BLOCKING_CALLS:
+        if _IN_TESTS and blocking_call.skip_for_tests:
+            continue
 
-    # Prevent sleeping in event loop.
-    time.sleep = protect_loop(
-        time.sleep,
-        check_allowed=_check_sleep_call_allowed,
-        loop_thread_id=loop_thread_id,
-    )
-
-    glob.glob = protect_loop(
-        glob.glob, strict_core=False, strict=False, loop_thread_id=loop_thread_id
-    )
-    glob.iglob = protect_loop(
-        glob.iglob, strict_core=False, strict=False, loop_thread_id=loop_thread_id
-    )
-    os.walk = protect_loop(
-        os.walk, strict_core=False, strict=False, loop_thread_id=loop_thread_id
-    )
-
-    if not _IN_TESTS:
-        # Prevent files being opened inside the event loop
-        os.listdir = protect_loop(  # type: ignore[assignment]
-            os.listdir, strict_core=False, strict=False, loop_thread_id=loop_thread_id
-        )
-        os.scandir = protect_loop(  # type: ignore[assignment]
-            os.scandir, strict_core=False, strict=False, loop_thread_id=loop_thread_id
-        )
-
-        builtins.open = protect_loop(  # type: ignore[assignment]
-            builtins.open,
-            strict_core=False,
-            strict=False,
-            check_allowed=_check_file_allowed,
-            loop_thread_id=loop_thread_id,
-        )
-        # unittest uses `importlib.import_module` to do mocking
-        # so we cannot protect it if we are running tests
-        importlib.import_module = protect_loop(
-            importlib.import_module,
-            strict_core=False,
-            strict=False,
-            check_allowed=_check_import_call_allowed,
+        protected_function = protect_loop(
+            blocking_call.original_func,
+            strict=blocking_call.strict,
+            strict_core=blocking_call.strict_core,
+            check_allowed=blocking_call.check_allowed,
             loop_thread_id=loop_thread_id,
         )
+        setattr(blocking_call.object, blocking_call.function, protected_function)
+        calls.add(blocking_call)