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
    pr.pause_writing()
    assert pr._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_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 not pr._connection_lost
    pr.connection_lost(None)
    assert pr.transport is None
    assert pr._connection_lost


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 not pr._connection_lost
    pr.pause_writing()
    pr.connection_lost(None)
    assert pr.transport is None
    assert pr._connection_lost


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()
    assert (await t) is None
    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)
    assert (await t) is None
    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(RuntimeError) as cm:
        await t
    assert cm.value 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():
        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():
        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

[((180, 204), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ([], {}), '()\n', (202, 204), False, 'import asyncio\n'), ((209, 237), 'asyncio.set_event_loop', 'asyncio.set_event_loop', (['None'], {}), '(None)\n', (231, 237), False, 'import asyncio\n'), ((247, 265), 'aiohttp.base_protocol.BaseProtocol', 'BaseProtocol', (['loop'], {}), '(loop)\n', (259, 265), False, 'from aiohttp.base_protocol import BaseProtocol\n'), ((347, 371), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ([], {}), '()\n', (369, 371), False, 'import asyncio\n'), ((381, 399), 'aiohttp.base_protocol.BaseProtocol', 'BaseProtocol', (['loop'], {}), '(loop)\n', (393, 399), False, 'from aiohttp.base_protocol import BaseProtocol\n'), ((536, 560), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ([], {}), '()\n', (558, 560), False, 'import asyncio\n'), ((570, 593), 'aiohttp.base_protocol.BaseProtocol', 'BaseProtocol', ([], {'loop': 'loop'}), '(loop=loop)\n', (582, 593), False, 'from aiohttp.base_protocol import BaseProtocol\n'), ((744, 768), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ([], {}), '()\n', (766, 768), False, 'import asyncio\n'), ((778, 801), 'aiohttp.base_protocol.BaseProtocol', 'BaseProtocol', ([], {'loop': 'loop'}), '(loop=loop)\n', (790, 801), False, 'from aiohttp.base_protocol import BaseProtocol\n'), ((811, 822), 'unittest.mock.Mock', 'mock.Mock', ([], {}), '()\n', (820, 822), False, 'from unittest import mock\n'), ((984, 1008), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ([], {}), '()\n', (1006, 1008), False, 'import asyncio\n'), ((1018, 1041), 'aiohttp.base_protocol.BaseProtocol', 'BaseProtocol', ([], {'loop': 'loop'}), '(loop=loop)\n', (1030, 1041), False, 'from aiohttp.base_protocol import BaseProtocol\n'), ((1051, 1062), 'unittest.mock.Mock', 'mock.Mock', ([], {}), '()\n', (1060, 1062), False, 'from unittest import mock\n'), ((1294, 1318), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ([], {}), '()\n', (1316, 1318), False, 'import asyncio\n'), ((1328, 1351), 'aiohttp.base_protocol.BaseProtocol', 'BaseProtocol', ([], {'loop': 'loop'}), '(loop=loop)\n', (1340, 1351), False, 'from aiohttp.base_protocol import BaseProtocol\n'), ((1361, 1372), 'unittest.mock.Mock', 'mock.Mock', ([], {}), '()\n', (1370, 1372), False, 'from unittest import mock\n'), ((1600, 1624), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ([], {}), '()\n', (1622, 1624), False, 'import asyncio\n'), ((1634, 1657), 'aiohttp.base_protocol.BaseProtocol', 'BaseProtocol', ([], {'loop': 'loop'}), '(loop=loop)\n', (1646, 1657), False, 'from aiohttp.base_protocol import BaseProtocol\n'), ((1667, 1678), 'unittest.mock.Mock', 'mock.Mock', ([], {}), '()\n', (1676, 1678), False, 'from unittest import mock\n'), ((1870, 1894), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ([], {}), '()\n', (1892, 1894), False, 'import asyncio\n'), ((1904, 1927), 'aiohttp.base_protocol.BaseProtocol', 'BaseProtocol', ([], {'loop': 'loop'}), '(loop=loop)\n', (1916, 1927), False, 'from aiohttp.base_protocol import BaseProtocol\n'), ((1937, 1948), 'unittest.mock.Mock', 'mock.Mock', ([], {}), '()\n', (1946, 1948), False, 'from unittest import mock\n'), ((2136, 2160), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ([], {}), '()\n', (2158, 2160), False, 'import asyncio\n'), ((2170, 2193), 'aiohttp.base_protocol.BaseProtocol', 'BaseProtocol', ([], {'loop': 'loop'}), '(loop=loop)\n', (2182, 2193), False, 'from aiohttp.base_protocol import BaseProtocol\n'), ((2203, 2214), 'unittest.mock.Mock', 'mock.Mock', ([], {}), '()\n', (2212, 2214), False, 'from unittest import mock\n'), ((2528, 2552), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ([], {}), '()\n', (2550, 2552), False, 'import asyncio\n'), ((2562, 2585), 'aiohttp.base_protocol.BaseProtocol', 'BaseProtocol', ([], {'loop': 'loop'}), '(loop=loop)\n', (2574, 2585), False, 'from aiohttp.base_protocol import BaseProtocol\n'), ((2595, 2606), 'unittest.mock.Mock', 'mock.Mock', ([], {}), '()\n', (2604, 2606), False, 'from unittest import mock\n'), ((2932, 2956), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ([], {}), '()\n', (2954, 2956), False, 'import asyncio\n'), ((2966, 2989), 'aiohttp.base_protocol.BaseProtocol', 'BaseProtocol', ([], {'loop': 'loop'}), '(loop=loop)\n', (2978, 2989), False, 'from aiohttp.base_protocol import BaseProtocol\n'), ((2999, 3010), 'unittest.mock.Mock', 'mock.Mock', ([], {}), '()\n', (3008, 3010), False, 'from unittest import mock\n'), ((3411, 3435), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ([], {}), '()\n', (3433, 3435), False, 'import asyncio\n'), ((3445, 3468), 'aiohttp.base_protocol.BaseProtocol', 'BaseProtocol', ([], {'loop': 'loop'}), '(loop=loop)\n', (3457, 3468), False, 'from aiohttp.base_protocol import BaseProtocol\n'), ((3478, 3489), 'unittest.mock.Mock', 'mock.Mock', ([], {}), '()\n', (3487, 3489), False, 'from unittest import mock\n'), ((3947, 3971), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ([], {}), '()\n', (3969, 3971), False, 'import asyncio\n'), ((3981, 4004), 'aiohttp.base_protocol.BaseProtocol', 'BaseProtocol', ([], {'loop': 'loop'}), '(loop=loop)\n', (3993, 4004), False, 'from aiohttp.base_protocol import BaseProtocol\n'), ((4014, 4025), 'unittest.mock.Mock', 'mock.Mock', ([], {}), '()\n', (4023, 4025), False, 'from unittest import mock\n'), ((1744, 1779), 'pytest.raises', 'pytest.raises', (['ConnectionResetError'], {}), '(ConnectionResetError)\n', (1757, 1779), False, 'import pytest\n'), ((2321, 2337), 'asyncio.sleep', 'asyncio.sleep', (['(0)'], {}), '(0)\n', (2334, 2337), False, 'import asyncio\n'), ((2713, 2729), 'asyncio.sleep', 'asyncio.sleep', (['(0)'], {}), '(0)\n', (2726, 2729), False, 'import asyncio\n'), ((3117, 3133), 'asyncio.sleep', 'asyncio.sleep', (['(0)'], {}), '(0)\n', (3130, 3133), False, 'import asyncio\n'), ((3237, 3264), 'pytest.raises', 'pytest.raises', (['RuntimeError'], {}), '(RuntimeError)\n', (3250, 3264), False, 'import pytest\n'), ((3799, 3831), 'contextlib.suppress', 'suppress', (['asyncio.CancelledError'], {}), '(asyncio.CancelledError)\n', (3807, 3831), False, 'from contextlib import suppress\n'), ((4330, 4362), 'contextlib.suppress', 'suppress', (['asyncio.CancelledError'], {}), '(asyncio.CancelledError)\n', (4338, 4362), False, 'from contextlib import suppress\n')]