Repository: aaugustin/websockets
Pull request #1307 (2023-03-07 21:56:41)

Switch to using asyncio.timeout instead of asyncio.wait_for
`asyncio.wait_for` creates a task whereas `asyncio.timeout` avoids doing this which decreases latency.

Fallback to using `async_timeout` when the python version is too old (<3.11)

`asyncio.timeout` will become the underlying implementation for `async.wait_for` in cpython 3.12
https://github.com/python/cpython/pull/98518

aaugustin: Thank you for the suggestion! I think it's a good idea. Let's do it.

Until now, websockets has no dependencies. This removes an entire category of concerns like managing versions that, as a single maintainer, I don't want to deal with. I understand that this may not look like a big benefit but it is valuable to me.

In order to keep this benefit, I am tempted to do one of the following:

1. create a small compatibility layer to use asyncio.wait_for on Python < 3.11 and asyncio.timeout on Python ≥ 3.12
2. vendor CPython's or  async-timeout's implementation — they're probably identical if async-timeout was merged into CPython.

Option 1 could be as simple as:

```python
if sys.version_info[:2] >= (3, 12):

    async def asyncio_wait_for(coro, timeout):
        async with asyncio.timeout(timeout):
            await coro

else:

    asyncio_wait_for = asyncio.wait_for
```

In all cases, version compatibility code should go in [websockets.legacy.compatibility](https://github.com/aaugustin/websockets/blob/main/src/websockets/legacy/compatibility.py) rather than a new module.
aaugustin: I'm aware that my option 1 doesn't yield performance benefits on Python < 3.12 while your version does. IMO those who care about performance optimizations available in Python 3.12 should upgrade to Python 3.12. It isn't my job to fake that for them.
bdraco: I'll work on option 2 since the races in `async.wait_for` and the issue tickets they generated for `cpython` were part of the motivation for cpython switching out the implementation.  It think it would be best if everything behaved the same regardless of the python version.

ex:
https://github.com/python/cpython/issues/86296
https://github.com/python/cpython/issues/81839
https://github.com/python/cpython/issues/96764
bdraco: Its passing locally now. Hopefully it should be good to go now
aaugustin: Also, It would be great if you could rebase and squash into a final commit to clean up the PR.
bdraco: Late here.  I'll take a look in the morning if the Home Assistant beta isn't too busy.
aaugustin: I squashed your commits and I'm finalizing the changes here: #1325.

Patch:
diff --git a/docs/faq/common.rst b/docs/faq/common.rst
index 2a512ea..505149a 100644
--- a/docs/faq/common.rst
+++ b/docs/faq/common.rst
@@ -105,9 +105,14 @@ you can adjust it with the ``ping_timeout`` argument.
 How do I set a timeout on :meth:`~legacy.protocol.WebSocketCommonProtocol.recv`?
 --------------------------------------------------------------------------------
 
-Use :func:`~asyncio.wait_for`::
+On Python ≥ 3.11, use :func:`asyncio.timeout`::
 
-    await asyncio.wait_for(websocket.recv(), timeout=10)
+    async with asyncio.timeout(timeout=10):
+        message = await websocket.recv()
+
+On older versions of Python, use :func:`asyncio.wait_for`::
+
+    message = await asyncio.wait_for(websocket.recv(), timeout=10)
 
 This technique works for most APIs. When it doesn't, for example with
 asynchronous context managers, websockets provides an ``open_timeout`` argument.
diff --git a/docs/topics/design.rst b/docs/topics/design.rst
index 33dd187..f164d29 100644
--- a/docs/topics/design.rst
+++ b/docs/topics/design.rst
@@ -437,7 +437,7 @@ propagate cancellation to them.
 prevent cancellation.
 
 :meth:`~legacy.protocol.WebSocketCommonProtocol.close` waits for the data transfer
-task to terminate with :func:`~asyncio.wait_for`. If it's canceled or if the
+task to terminate with :func:`~asyncio.timeout`. If it's canceled or if the
 timeout elapses, :attr:`~legacy.protocol.WebSocketCommonProtocol.transfer_data_task`
 is canceled, which is correct at this point.
 :meth:`~legacy.protocol.WebSocketCommonProtocol.close` then waits for
diff --git a/pyproject.toml b/pyproject.toml
index 989b6b5..0707c64 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,6 +39,7 @@ branch = true
 omit = [
     # */websockets matches src/websockets and .tox/**/site-packages/websockets
     "*/websockets/__main__.py",
+    "*/websockets/legacy/async_timeout.py",
     "*/websockets/legacy/compatibility.py",
     "tests/maxi_cov.py",
 ]
diff --git a/src/websockets/legacy/async_timeout.py b/src/websockets/legacy/async_timeout.py
new file mode 100644
index 0000000..8264094
--- /dev/null
+++ b/src/websockets/legacy/async_timeout.py
@@ -0,0 +1,265 @@
+# From https://github.com/aio-libs/async-timeout/blob/master/async_timeout/__init__.py
+# Licensed under the Apache License (Apache-2.0)
+
+import asyncio
+import enum
+import sys
+import warnings
+from types import TracebackType
+from typing import Optional, Type
+
+
+# From https://github.com/python/typing_extensions/blob/main/src/typing_extensions.py
+# Licensed under the Python Software Foundation License (PSF-2.0)
+
+if sys.version_info >= (3, 11):
+    from typing import final
+else:
+    # @final exists in 3.8+, but we backport it for all versions
+    # before 3.11 to keep support for the __final__ attribute.
+    # See https://bugs.python.org/issue46342
+    def final(f):
+        """This decorator can be used to indicate to type checkers that
+        the decorated method cannot be overridden, and decorated class
+        cannot be subclassed. For example:
+
+            class Base:
+                @final
+                def done(self) -> None:
+                    ...
+            class Sub(Base):
+                def done(self) -> None:  # Error reported by type checker
+                    ...
+            @final
+            class Leaf:
+                ...
+            class Other(Leaf):  # Error reported by type checker
+                ...
+
+        There is no runtime checking of these properties. The decorator
+        sets the ``__final__`` attribute to ``True`` on the decorated object
+        to allow runtime introspection.
+        """
+        try:
+            f.__final__ = True
+        except (AttributeError, TypeError):
+            # Skip the attribute silently if it is not writable.
+            # AttributeError happens if the object has __slots__ or a
+            # read-only property, TypeError if it's a builtin class.
+            pass
+        return f
+
+
+# End https://github.com/aio-libs/async-timeout/blob/master/async_timeout/__init__.py
+
+__version__ = "4.0.2"
+
+
+__all__ = ("timeout", "timeout_at", "Timeout")
+
+
+def timeout(delay: Optional[float]) -> "Timeout":
+    """timeout context manager.
+
+    Useful in cases when you want to apply timeout logic around block
+    of code or in cases when asyncio.wait_for is not suitable. For example:
+
+    >>> async with timeout(0.001):
+    ...     async with aiohttp.get('https://github.com') as r:
+    ...         await r.text()
+
+
+    delay - value in seconds or None to disable timeout logic
+    """
+    loop = asyncio.get_running_loop()
+    if delay is not None:
+        deadline = loop.time() + delay  # type: Optional[float]
+    else:
+        deadline = None
+    return Timeout(deadline, loop)
+
+
+def timeout_at(deadline: Optional[float]) -> "Timeout":
+    """Schedule the timeout at absolute time.
+
+    deadline argument points on the time in the same clock system
+    as loop.time().
+
+    Please note: it is not POSIX time but a time with
+    undefined starting base, e.g. the time of the system power on.
+
+    >>> async with timeout_at(loop.time() + 10):
+    ...     async with aiohttp.get('https://github.com') as r:
+    ...         await r.text()
+
+
+    """
+    loop = asyncio.get_running_loop()
+    return Timeout(deadline, loop)
+
+
+class _State(enum.Enum):
+    INIT = "INIT"
+    ENTER = "ENTER"
+    TIMEOUT = "TIMEOUT"
+    EXIT = "EXIT"
+
+
+@final
+class Timeout:
+    # Internal class, please don't instantiate it directly
+    # Use timeout() and timeout_at() public factories instead.
+    #
+    # Implementation note: `async with timeout()` is preferred
+    # over `with timeout()`.
+    # While technically the Timeout class implementation
+    # doesn't need to be async at all,
+    # the `async with` statement explicitly points that
+    # the context manager should be used from async function context.
+    #
+    # This design allows to avoid many silly misusages.
+    #
+    # TimeoutError is raised immediately when scheduled
+    # if the deadline is passed.
+    # The purpose is to time out as soon as possible
+    # without waiting for the next await expression.
+
+    __slots__ = ("_deadline", "_loop", "_state", "_timeout_handler")
+
+    def __init__(
+        self, deadline: Optional[float], loop: asyncio.AbstractEventLoop
+    ) -> None:
+        self._loop = loop
+        self._state = _State.INIT
+
+        self._timeout_handler = None  # type: Optional[asyncio.Handle]
+        if deadline is None:
+            self._deadline = None  # type: Optional[float]
+        else:
+            self.update(deadline)
+
+    def __enter__(self) -> "Timeout":
+        warnings.warn(
+            "with timeout() is deprecated, use async with timeout() instead",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        self._do_enter()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> Optional[bool]:
+        self._do_exit(exc_type)
+        return None
+
+    async def __aenter__(self) -> "Timeout":
+        self._do_enter()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> Optional[bool]:
+        self._do_exit(exc_type)
+        return None
+
+    @property
+    def expired(self) -> bool:
+        """Is timeout expired during execution?"""
+        return self._state == _State.TIMEOUT
+
+    @property
+    def deadline(self) -> Optional[float]:
+        return self._deadline
+
+    def reject(self) -> None:
+        """Reject scheduled timeout if any."""
+        # cancel is maybe better name but
+        # task.cancel() raises CancelledError in asyncio world.
+        if self._state not in (_State.INIT, _State.ENTER):
+            raise RuntimeError(f"invalid state {self._state.value}")
+        self._reject()
+
+    def _reject(self) -> None:
+        if self._timeout_handler is not None:
+            self._timeout_handler.cancel()
+            self._timeout_handler = None
+
+    def shift(self, delay: float) -> None:
+        """Advance timeout on delay seconds.
+
+        The delay can be negative.
+
+        Raise RuntimeError if shift is called when deadline is not scheduled
+        """
+        deadline = self._deadline
+        if deadline is None:
+            raise RuntimeError("cannot shift timeout if deadline is not scheduled")
+        self.update(deadline + delay)
+
+    def update(self, deadline: float) -> None:
+        """Set deadline to absolute value.
+
+        deadline argument points on the time in the same clock system
+        as loop.time().
+
+        If new deadline is in the past the timeout is raised immediately.
+
+        Please note: it is not POSIX time but a time with
+        undefined starting base, e.g. the time of the system power on.
+        """
+        if self._state == _State.EXIT:
+            raise RuntimeError("cannot reschedule after exit from context manager")
+        if self._state == _State.TIMEOUT:
+            raise RuntimeError("cannot reschedule expired timeout")
+        if self._timeout_handler is not None:
+            self._timeout_handler.cancel()
+        self._deadline = deadline
+        if self._state != _State.INIT:
+            self._reschedule()
+
+    def _reschedule(self) -> None:
+        assert self._state == _State.ENTER
+        deadline = self._deadline
+        if deadline is None:
+            return
+
+        now = self._loop.time()
+        if self._timeout_handler is not None:
+            self._timeout_handler.cancel()
+
+        task = asyncio.current_task()
+        if deadline <= now:
+            self._timeout_handler = self._loop.call_soon(self._on_timeout, task)
+        else:
+            self._timeout_handler = self._loop.call_at(deadline, self._on_timeout, task)
+
+    def _do_enter(self) -> None:
+        if self._state != _State.INIT:
+            raise RuntimeError(f"invalid state {self._state.value}")
+        self._state = _State.ENTER
+        self._reschedule()
+
+    def _do_exit(self, exc_type: Optional[Type[BaseException]]) -> None:
+        if exc_type is asyncio.CancelledError and self._state == _State.TIMEOUT:
+            self._timeout_handler = None
+            raise asyncio.TimeoutError
+        # timeout has not expired
+        self._state = _State.EXIT
+        self._reject()
+        return None
+
+    def _on_timeout(self, task: "asyncio.Task[None]") -> None:
+        task.cancel()
+        self._state = _State.TIMEOUT
+        # drop the reference early
+        self._timeout_handler = None
+
+
+# End https://github.com/aio-libs/async-timeout/blob/master/async_timeout/__init__.py
diff --git a/src/websockets/legacy/client.py b/src/websockets/legacy/client.py
index b79bbab..c5e9d0d 100644
--- a/src/websockets/legacy/client.py
+++ b/src/websockets/legacy/client.py
@@ -44,6 +44,7 @@ from ..headers import (
 from ..http import USER_AGENT
 from ..typing import ExtensionHeader, LoggerLike, Origin, Subprotocol
 from ..uri import WebSocketURI, parse_uri
+from .compatibility import asyncio_timeout
 from .handshake import build_request, check_response
 from .http import read_response
 from .protocol import WebSocketCommonProtocol
@@ -650,7 +651,8 @@ class Connect:
         return self.__await_impl_timeout__().__await__()
 
     async def __await_impl_timeout__(self) -> WebSocketClientProtocol:
-        return await asyncio.wait_for(self.__await_impl__(), self.open_timeout)
+        async with asyncio_timeout(self.open_timeout):
+            return await self.__await_impl__()
 
     async def __await_impl__(self) -> WebSocketClientProtocol:
         for redirects in range(self.MAX_REDIRECTS_ALLOWED):
diff --git a/src/websockets/legacy/compatibility.py b/src/websockets/legacy/compatibility.py
index 303e203..cb9b02c 100644
--- a/src/websockets/legacy/compatibility.py
+++ b/src/websockets/legacy/compatibility.py
@@ -5,6 +5,9 @@ import sys
 from typing import Any, Dict
 
 
+__all__ = ["asyncio_timeout", "loop_if_py_lt_38"]
+
+
 if sys.version_info[:2] >= (3, 8):
 
     def loop_if_py_lt_38(loop: asyncio.AbstractEventLoop) -> Dict[str, Any]:
@@ -22,3 +25,9 @@ else:
 
         """
         return {"loop": loop}
+
+
+if sys.version_info[:2] >= (3, 11):
+    from asyncio import timeout as asyncio_timeout  # noqa: F401
+else:
+    from .async_timeout import timeout as asyncio_timeout  # noqa: F401
diff --git a/src/websockets/legacy/protocol.py b/src/websockets/legacy/protocol.py
index 7f9ab2b..733abb3 100644
--- a/src/websockets/legacy/protocol.py
+++ b/src/websockets/legacy/protocol.py
@@ -53,7 +53,7 @@ from ..frames import (
 )
 from ..protocol import State
 from ..typing import Data, LoggerLike, Subprotocol
-from .compatibility import loop_if_py_lt_38
+from .compatibility import asyncio_timeout, loop_if_py_lt_38
 from .framing import Frame
 
 
@@ -107,7 +107,8 @@ class WebSocketCommonProtocol(asyncio.Protocol):
       context manager;
     * on the server side, when the connection handler terminates.
 
-    To apply a timeout to any other API, wrap it in :func:`~asyncio.wait_for`.
+    To apply a timeout to any other API, wrap it in :func:`~asyncio.timeout` or
+    :func:`~asyncio.wait_for`.
 
     The ``max_size`` parameter enforces the maximum size for incoming messages
     in bytes. The default value is 1 MiB. If a larger message is received,
@@ -513,7 +514,7 @@ class WebSocketCommonProtocol(asyncio.Protocol):
         message. The next invocation of :meth:`recv` will return it.
 
         This makes it possible to enforce a timeout by wrapping :meth:`recv` in
-        :func:`~asyncio.wait_for`.
+        :func:`~asyncio.timeout` or :func:`~asyncio.wait_for`.
 
         Returns:
             Data: A string (:class:`str`) for a Text_ frame. A bytestring
@@ -761,19 +762,16 @@ class WebSocketCommonProtocol(asyncio.Protocol):
 
         """
         try:
-            await asyncio.wait_for(
-                self.write_close_frame(Close(code, reason)),
-                self.close_timeout,
-                **loop_if_py_lt_38(self.loop),
-            )
+            async with asyncio_timeout(self.close_timeout):
+                await self.write_close_frame(Close(code, reason))
         except asyncio.TimeoutError:
             # If the close frame cannot be sent because the send buffers
             # are full, the closing handshake won't complete anyway.
             # Fail the connection to shut down faster.
             self.fail_connection()
 
-        # If no close frame is received within the timeout, wait_for() cancels
-        # the data transfer task and raises TimeoutError.
+        # If no close frame is received within the timeout, asyncio_timeout()
+        # cancels the data transfer task and raises TimeoutError.
 
         # If close() is called multiple times concurrently and one of these
         # calls hits the timeout, the data transfer task will be canceled.
@@ -782,11 +780,8 @@ class WebSocketCommonProtocol(asyncio.Protocol):
         try:
             # If close() is canceled during the wait, self.transfer_data_task
             # is canceled before the timeout elapses.
-            await asyncio.wait_for(
-                self.transfer_data_task,
-                self.close_timeout,
-                **loop_if_py_lt_38(self.loop),
-            )
+            async with asyncio_timeout(self.close_timeout):
+                await self.transfer_data_task
         except (asyncio.TimeoutError, asyncio.CancelledError):
             pass
 
@@ -1268,11 +1263,8 @@ class WebSocketCommonProtocol(asyncio.Protocol):
 
                 if self.ping_timeout is not None:
                     try:
-                        await asyncio.wait_for(
-                            pong_waiter,
-                            self.ping_timeout,
-                            **loop_if_py_lt_38(self.loop),
-                        )
+                        async with asyncio_timeout(self.ping_timeout):
+                            await pong_waiter
                         self.logger.debug("% received keepalive pong")
                     except asyncio.TimeoutError:
                         if self.debug:
@@ -1369,7 +1361,8 @@ class WebSocketCommonProtocol(asyncio.Protocol):
         # Abort the TCP connection. Buffers are discarded.
         if self.debug:
             self.logger.debug("x aborting TCP connection")
-        self.transport.abort()
+        # Due to a bug in coverage, this is erroneously reported as not covered.
+        self.transport.abort()  # pragma: no cover
 
         # connection_lost() is called quickly after aborting.
         await self.wait_for_connection_lost()
@@ -1384,11 +1377,8 @@ class WebSocketCommonProtocol(asyncio.Protocol):
         """
         if not self.connection_lost_waiter.done():
             try:
-                await asyncio.wait_for(
-                    asyncio.shield(self.connection_lost_waiter),
-                    self.close_timeout,
-                    **loop_if_py_lt_38(self.loop),
-                )
+                async with asyncio_timeout(self.close_timeout):
+                    await asyncio.shield(self.connection_lost_waiter)
             except asyncio.TimeoutError:
                 pass
         # Re-check self.connection_lost_waiter.done() synchronously because
diff --git a/src/websockets/legacy/server.py b/src/websockets/legacy/server.py
index 3506276..25d5a71 100644
--- a/src/websockets/legacy/server.py
+++ b/src/websockets/legacy/server.py
@@ -46,7 +46,7 @@ from ..headers import (
 from ..http import USER_AGENT
 from ..protocol import State
 from ..typing import ExtensionHeader, LoggerLike, Origin, Subprotocol
-from .compatibility import loop_if_py_lt_38
+from .compatibility import asyncio_timeout, loop_if_py_lt_38
 from .handshake import build_response, check_request
 from .http import read_request
 from .protocol import WebSocketCommonProtocol
@@ -163,15 +163,13 @@ class WebSocketServerProtocol(WebSocketCommonProtocol):
         """
         try:
             try:
-                await asyncio.wait_for(
-                    self.handshake(
+                async with asyncio_timeout(self.open_timeout):
+                    await self.handshake(
                         origins=self.origins,
                         available_extensions=self.available_extensions,
                         available_subprotocols=self.available_subprotocols,
                         extra_headers=self.extra_headers,
-                    ),
-                    self.open_timeout,
-                )
+                    )
             # Remove this branch when dropping support for Python < 3.8
             # because CancelledError no longer inherits Exception.
             except asyncio.CancelledError:  # pragma: no cover


Test patch:
diff --git a/tests/legacy/test_client_server.py b/tests/legacy/test_client_server.py
index 3dd06d0..133af05 100644
--- a/tests/legacy/test_client_server.py
+++ b/tests/legacy/test_client_server.py
@@ -29,6 +29,7 @@ from websockets.extensions.permessage_deflate import (
 )
 from websockets.http import USER_AGENT
 from websockets.legacy.client import *
+from websockets.legacy.compatibility import asyncio_timeout
 from websockets.legacy.handshake import build_response
 from websockets.legacy.http import read_response
 from websockets.legacy.server import *
@@ -1129,7 +1130,8 @@ class CommonClientServerTests:
 
         async def cancelled_client():
             start_client = connect(get_server_uri(self.server), sock=sock)
-            await asyncio.wait_for(start_client, 5 * MS)
+            async with asyncio_timeout(5 * MS):
+                await start_client
 
         with self.assertRaises(asyncio.TimeoutError):
             self.loop.run_until_complete(cancelled_client())
diff --git a/tests/legacy/test_protocol.py b/tests/legacy/test_protocol.py
index ab8155b..a05dcc6 100644
--- a/tests/legacy/test_protocol.py
+++ b/tests/legacy/test_protocol.py
@@ -190,7 +190,6 @@ class CommonTests:
         close_frame_data = Close(code, reason).serialize()
         # Trigger the closing handshake from the local endpoint.
         close_task = self.loop.create_task(self.protocol.close(code, reason))
-        self.run_loop_once()  # wait_for executes
         self.run_loop_once()  # write_frame executes
         # Empty the outgoing data stream so we can make assertions later on.
         self.assertOneFrameSent(True, OP_CLOSE, close_frame_data)
@@ -919,6 +918,7 @@ class CommonTests:
         close_task = self.half_close_connection_local()
 
         self.receive_frame(Frame(True, OP_PING, b"test"))
+        self.run_loop_once()
 
         with self.assertNoLogs():
             self.loop.run_until_complete(self.protocol.close())
@@ -931,6 +931,7 @@ class CommonTests:
         # which prevents responding with a pong frame properly.
         self.receive_frame(Frame(True, OP_PING, b"test"))
         self.receive_eof()
+        self.run_loop_once()
 
         with self.assertNoLogs():
             self.loop.run_until_complete(self.protocol.close())
@@ -1362,6 +1363,7 @@ class CommonTests:
         # which prevents echoing the close frame properly.
         self.receive_frame(self.close_frame)
         self.receive_eof()
+        self.run_loop_once()
 
         with self.assertNoLogs():
             self.loop.run_until_complete(self.protocol.close(reason="oh noes!"))
@@ -1375,6 +1377,7 @@ class CommonTests:
         # https://github.com/aaugustin/websockets/issues/339
         self.loop.call_soon(self.receive_frame, self.remote_close)
         self.loop.call_soon(self.receive_eof_if_client)
+        self.run_loop_once()
 
         self.loop.run_until_complete(self.protocol.close(reason="local"))
 
@@ -1386,6 +1389,7 @@ class CommonTests:
 
     def test_close_preserves_incoming_frames(self):
         self.receive_frame(Frame(True, OP_TEXT, b"hello"))
+        self.run_loop_once()
 
         self.loop.call_later(MS, self.receive_frame, self.close_frame)
         self.loop.call_later(MS, self.receive_eof_if_client)
@@ -1573,8 +1577,10 @@ class ServerTests(CommonTests, AsyncioTestCase):
             # HACK: disable write_eof => other end drops connection emulation.
             self.transport._eof = True
             self.receive_frame(self.close_frame)
+            self.run_loop_once()
             self.loop.run_until_complete(self.protocol.close(reason="close"))
-        self.assertConnectionClosed(1000, "close")
+        # Due to a bug in coverage, this is erroneously reported as not covered.
+        self.assertConnectionClosed(1000, "close")  # pragma: no cover
 
     def test_local_close_connection_lost_timeout_after_close(self):
         self.protocol.close_timeout = 10 * MS
@@ -1589,8 +1595,10 @@ class ServerTests(CommonTests, AsyncioTestCase):
             # HACK: disable close => other end drops connection emulation.
             self.transport._closing = True
             self.receive_frame(self.close_frame)
+            self.run_loop_once()
             self.loop.run_until_complete(self.protocol.close(reason="close"))
-        self.assertConnectionClosed(1000, "close")
+        # Due to a bug in coverage, this is erroneously reported as not covered.
+        self.assertConnectionClosed(1000, "close")  # pragma: no cover
 
 
 class ClientTests(CommonTests, AsyncioTestCase):
@@ -1608,7 +1616,8 @@ class ClientTests(CommonTests, AsyncioTestCase):
         # Check the timing within -1/+9ms for robustness.
         with self.assertCompletesWithin(19 * MS, 29 * MS):
             self.loop.run_until_complete(self.protocol.close(reason="close"))
-        self.assertConnectionClosed(1006, "")
+        # Due to a bug in coverage, this is erroneously reported as not covered.
+        self.assertConnectionClosed(1006, "")  # pragma: no cover
 
     def test_local_close_receive_close_frame_timeout(self):
         self.protocol.close_timeout = 10 * MS
@@ -1618,7 +1627,8 @@ class ClientTests(CommonTests, AsyncioTestCase):
         # Check the timing within -1/+9ms for robustness.
         with self.assertCompletesWithin(19 * MS, 29 * MS):
             self.loop.run_until_complete(self.protocol.close(reason="close"))
-        self.assertConnectionClosed(1006, "")
+        # Due to a bug in coverage, this is erroneously reported as not covered.
+        self.assertConnectionClosed(1006, "")  # pragma: no cover
 
     def test_local_close_connection_lost_timeout_after_write_eof(self):
         self.protocol.close_timeout = 10 * MS
@@ -1631,8 +1641,10 @@ class ClientTests(CommonTests, AsyncioTestCase):
             # HACK: disable write_eof => other end drops connection emulation.
             self.transport._eof = True
             self.receive_frame(self.close_frame)
+            self.run_loop_once()
             self.loop.run_until_complete(self.protocol.close(reason="close"))
-        self.assertConnectionClosed(1000, "close")
+        # Due to a bug in coverage, this is erroneously reported as not covered.
+        self.assertConnectionClosed(1000, "close")  # pragma: no cover
 
     def test_local_close_connection_lost_timeout_after_close(self):
         self.protocol.close_timeout = 10 * MS
@@ -1650,5 +1662,7 @@ class ClientTests(CommonTests, AsyncioTestCase):
             # HACK: disable close => other end drops connection emulation.
             self.transport._closing = True
             self.receive_frame(self.close_frame)
+            self.run_loop_once()
             self.loop.run_until_complete(self.protocol.close(reason="close"))
-        self.assertConnectionClosed(1000, "close")
+        # Due to a bug in coverage, this is erroneously reported as not covered.
+        self.assertConnectionClosed(1000, "close")  # pragma: no cover