test_base_protocol.py

import asyncio
from contextlib import suppress
from unittest import mock

import pytest

from aiohttp.base_protocol import BaseProtocol


async def test_loop() -> None:
    loop = asyncio.get_event_loop()
    asyncio.set_event_loop(None)
    pr = BaseProtocol(loop)
    assert pr._loop is loop


async def test_pause_writing() -> None:
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop)
    assert not pr._paused
    assert pr.writing_paused is False
    pr.pause_writing()
    assert pr._paused
    assert pr.writing_paused is True  # type: ignore[unreachable]


async def test_pause_reading_no_transport() -> None:
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop)
    assert not pr._reading_paused
    pr.pause_reading()
    assert not pr._reading_paused


async def test_pause_reading_stub_transport() -> None:
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop)
    tr = asyncio.Transport()
    pr.transport = tr
    assert not pr._reading_paused
    pr.pause_reading()
    assert pr._reading_paused


async def test_resume_reading_no_transport() -> None:
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop)
    pr._reading_paused = True
    pr.resume_reading()
    assert pr._reading_paused


async def test_resume_reading_stub_transport() -> None:
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop)
    tr = asyncio.Transport()
    pr.transport = tr
    pr._reading_paused = True
    pr.resume_reading()
    assert not pr._reading_paused


async def test_resume_writing_no_waiters() -> None:
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop=loop)
    pr.pause_writing()
    assert pr._paused
    pr.resume_writing()
    assert not pr._paused


async def test_resume_writing_waiter_done() -> None:
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop=loop)
    waiter = mock.Mock(done=mock.Mock(return_value=True))
    pr._drain_waiter = waiter
    pr._paused = True
    pr.resume_writing()
    assert not pr._paused
    assert waiter.mock_calls == [mock.call.done()]


async def test_connection_made() -> None:
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop=loop)
    tr = mock.Mock()
    assert pr.transport is None
    pr.connection_made(tr)
    assert pr.transport is not None


async def test_connection_lost_not_paused() -> None:
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop=loop)
    tr = mock.Mock()
    pr.connection_made(tr)
    assert pr.connected
    pr.connection_lost(None)
    assert pr.transport is None
    assert not pr.connected


async def test_connection_lost_paused_without_waiter() -> None:
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop=loop)
    tr = mock.Mock()
    pr.connection_made(tr)
    assert pr.connected
    pr.pause_writing()
    pr.connection_lost(None)
    assert pr.transport is None
    assert not pr.connected


async def test_connection_lost_waiter_done() -> None:
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop=loop)
    pr._paused = True
    waiter = mock.Mock(done=mock.Mock(return_value=True))
    pr._drain_waiter = waiter
    pr.connection_lost(None)
    assert pr._drain_waiter is None
    assert waiter.mock_calls == [mock.call.done()]  # type: ignore[unreachable]


async def test_drain_lost() -> None:
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop=loop)
    tr = mock.Mock()
    pr.connection_made(tr)
    pr.connection_lost(None)
    with pytest.raises(ConnectionResetError):
        await pr._drain_helper()


async def test_drain_not_paused() -> None:
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop=loop)
    tr = mock.Mock()
    pr.connection_made(tr)
    assert pr._drain_waiter is None
    await pr._drain_helper()
    assert pr._drain_waiter is None


async def test_resume_drain_waited() -> None:
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop=loop)
    tr = mock.Mock()
    pr.connection_made(tr)
    pr.pause_writing()

    t = loop.create_task(pr._drain_helper())
    await asyncio.sleep(0)

    assert pr._drain_waiter is not None
    pr.resume_writing()
    await t
    assert pr._drain_waiter is None


async def test_lost_drain_waited_ok() -> None:
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop=loop)
    tr = mock.Mock()
    pr.connection_made(tr)
    pr.pause_writing()

    t = loop.create_task(pr._drain_helper())
    await asyncio.sleep(0)

    assert pr._drain_waiter is not None
    pr.connection_lost(None)
    await t
    assert pr._drain_waiter is None


async def test_lost_drain_waited_exception() -> None:
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop=loop)
    tr = mock.Mock()
    pr.connection_made(tr)
    pr.pause_writing()

    t = loop.create_task(pr._drain_helper())
    await asyncio.sleep(0)

    assert pr._drain_waiter is not None
    exc = RuntimeError()
    pr.connection_lost(exc)
    with pytest.raises(ConnectionError, match=r"^Connection lost$") as cm:
        await t
    assert cm.value.__cause__ is exc
    assert pr._drain_waiter is None


async def test_lost_drain_cancelled() -> None:
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop=loop)
    tr = mock.Mock()
    pr.connection_made(tr)
    pr.pause_writing()

    fut = loop.create_future()

    async def wait() -> None:
        fut.set_result(None)
        await pr._drain_helper()

    t = loop.create_task(wait())
    await fut
    t.cancel()

    assert pr._drain_waiter is not None
    pr.connection_lost(None)
    with suppress(asyncio.CancelledError):
        await t
    assert pr._drain_waiter is None


async def test_resume_drain_cancelled() -> None:
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop=loop)
    tr = mock.Mock()
    pr.connection_made(tr)
    pr.pause_writing()

    fut = loop.create_future()

    async def wait() -> None:
        fut.set_result(None)
        await pr._drain_helper()

    t = loop.create_task(wait())
    await fut
    t.cancel()

    assert pr._drain_waiter is not None
    pr.resume_writing()
    with suppress(asyncio.CancelledError):
        await t
    assert pr._drain_waiter is None


async def test_cancelled_drain_no_unhandled_future_warning() -> None:
    """Cancelling a task during backpressure must not leave an orphaned future.

    When the handler task is cancelled while awaiting _drain_helper and
    connection_lost fires with an exception afterward, the waiter should
    already be done (cancelled) so set_exception is skipped. No "Future
    exception was never retrieved" warning should appear.

    Regression test for https://github.com/aio-libs/aiohttp/issues/12281
    """
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop=loop)
    tr = mock.Mock()
    pr.connection_made(tr)
    pr.pause_writing()

    fut = loop.create_future()

    async def wait() -> None:
        fut.set_result(None)
        await pr._drain_helper()

    t = loop.create_task(wait())
    await fut
    t.cancel()
    with suppress(asyncio.CancelledError):
        await t

    # After cancellation the waiter should be done (cancelled), so
    # connection_lost with an exception must not call set_exception.
    assert pr._drain_waiter is not None
    waiter = pr._drain_waiter
    assert waiter.done(), "waiter must be cancelled when task is cancelled"

    # This previously left an orphaned future with an unhandled exception
    # because asyncio.shield kept the original waiter alive and uncancelled.
    exc = RuntimeError("connection died")
    pr.connection_lost(exc)
    assert pr._drain_waiter is None

    # Verify the waiter is cancelled, not set with an exception.
    assert waiter.cancelled()  # type: ignore[unreachable]


async def test_parallel_drain_race_condition() -> None:
    loop = asyncio.get_event_loop()
    pr = BaseProtocol(loop=loop)
    tr = mock.Mock()
    pr.connection_made(tr)
    pr.pause_writing()

    ts = [loop.create_task(pr._drain_helper()) for _ in range(5)]
    assert not (await asyncio.wait(ts, timeout=0.5))[
        0
    ], "All draining tasks must be pending"

    assert pr._drain_waiter is not None
    pr.resume_writing()
    await asyncio.gather(*ts)
    assert pr._drain_waiter is None