cpython/Lib/test/test_asyncio/test_locks.py

"""Tests for locks.py"""

import unittest
from unittest import mock
import re

import asyncio
import collections

STR_RGX_REPR = (
    r'^<(?P<class>.*?) object at (?P<address>.*?)'
    r'\[(?P<extras>'
    r'(set|unset|locked|unlocked|filling|draining|resetting|broken)'
    r'(, value:\d)?'
    r'(, waiters:\d+)?'
    r'(, waiters:\d+\/\d+)?' # barrier
    r')\]>\Z'
)
RGX_REPR = re.compile(STR_RGX_REPR)


def tearDownModule():
    asyncio.set_event_loop_policy(None)


class LockTests(unittest.IsolatedAsyncioTestCase):

    async def test_repr(self):
        lock = asyncio.Lock()
        self.assertTrue(repr(lock).endswith('[unlocked]>'))
        self.assertTrue(RGX_REPR.match(repr(lock)))

        await lock.acquire()
        self.assertTrue(repr(lock).endswith('[locked]>'))
        self.assertTrue(RGX_REPR.match(repr(lock)))

    async def test_lock(self):
        lock = asyncio.Lock()

        with self.assertRaisesRegex(
            TypeError,
            "'Lock' object can't be awaited"
        ):
            await lock

        self.assertFalse(lock.locked())

    async def test_lock_doesnt_accept_loop_parameter(self):
        primitives_cls = [
            asyncio.Lock,
            asyncio.Condition,
            asyncio.Event,
            asyncio.Semaphore,
            asyncio.BoundedSemaphore,
        ]

        loop = asyncio.get_running_loop()

        for cls in primitives_cls:
            with self.assertRaisesRegex(
                TypeError,
                rf"{cls.__name__}\.__init__\(\) got an unexpected "
                rf"keyword argument 'loop'"
            ):
                cls(loop=loop)

    async def test_lock_by_with_statement(self):
        primitives = [
            asyncio.Lock(),
            asyncio.Condition(),
            asyncio.Semaphore(),
            asyncio.BoundedSemaphore(),
        ]

        for lock in primitives:
            await asyncio.sleep(0.01)
            self.assertFalse(lock.locked())
            with self.assertRaisesRegex(
                TypeError,
                r"'\w+' object can't be awaited"
            ):
                with await lock:
                    pass
            self.assertFalse(lock.locked())

    async def test_acquire(self):
        lock = asyncio.Lock()
        result = []

        self.assertTrue(await lock.acquire())

        async def c1(result):
            if await lock.acquire():
                result.append(1)
            return True

        async def c2(result):
            if await lock.acquire():
                result.append(2)
            return True

        async def c3(result):
            if await lock.acquire():
                result.append(3)
            return True

        t1 = asyncio.create_task(c1(result))
        t2 = asyncio.create_task(c2(result))

        await asyncio.sleep(0)
        self.assertEqual([], result)

        lock.release()
        await asyncio.sleep(0)
        self.assertEqual([1], result)

        await asyncio.sleep(0)
        self.assertEqual([1], result)

        t3 = asyncio.create_task(c3(result))

        lock.release()
        await asyncio.sleep(0)
        self.assertEqual([1, 2], result)

        lock.release()
        await asyncio.sleep(0)
        self.assertEqual([1, 2, 3], result)

        self.assertTrue(t1.done())
        self.assertTrue(t1.result())
        self.assertTrue(t2.done())
        self.assertTrue(t2.result())
        self.assertTrue(t3.done())
        self.assertTrue(t3.result())

    async def test_acquire_cancel(self):
        lock = asyncio.Lock()
        self.assertTrue(await lock.acquire())

        task = asyncio.create_task(lock.acquire())
        asyncio.get_running_loop().call_soon(task.cancel)
        with self.assertRaises(asyncio.CancelledError):
            await task
        self.assertFalse(lock._waiters)

    async def test_cancel_race(self):
        # Several tasks:
        # - A acquires the lock
        # - B is blocked in acquire()
        # - C is blocked in acquire()
        #
        # Now, concurrently:
        # - B is cancelled
        # - A releases the lock
        #
        # If B's waiter is marked cancelled but not yet removed from
        # _waiters, A's release() call will crash when trying to set
        # B's waiter; instead, it should move on to C's waiter.

        # Setup: A has the lock, b and c are waiting.
        lock = asyncio.Lock()

        async def lockit(name, blocker):
            await lock.acquire()
            try:
                if blocker is not None:
                    await blocker
            finally:
                lock.release()

        fa = asyncio.get_running_loop().create_future()
        ta = asyncio.create_task(lockit('A', fa))
        await asyncio.sleep(0)
        self.assertTrue(lock.locked())
        tb = asyncio.create_task(lockit('B', None))
        await asyncio.sleep(0)
        self.assertEqual(len(lock._waiters), 1)
        tc = asyncio.create_task(lockit('C', None))
        await asyncio.sleep(0)
        self.assertEqual(len(lock._waiters), 2)

        # Create the race and check.
        # Without the fix this failed at the last assert.
        fa.set_result(None)
        tb.cancel()
        self.assertTrue(lock._waiters[0].cancelled())
        await asyncio.sleep(0)
        self.assertFalse(lock.locked())
        self.assertTrue(ta.done())
        self.assertTrue(tb.cancelled())
        await tc

    async def test_cancel_release_race(self):
        # Issue 32734
        # Acquire 4 locks, cancel second, release first
        # and 2 locks are taken at once.
        loop = asyncio.get_running_loop()
        lock = asyncio.Lock()
        lock_count = 0
        call_count = 0

        async def lockit():
            nonlocal lock_count
            nonlocal call_count
            call_count += 1
            await lock.acquire()
            lock_count += 1

        def trigger():
            t1.cancel()
            lock.release()

        await lock.acquire()

        t1 = asyncio.create_task(lockit())
        t2 = asyncio.create_task(lockit())
        t3 = asyncio.create_task(lockit())

        # Start scheduled tasks
        await asyncio.sleep(0)

        loop.call_soon(trigger)
        with self.assertRaises(asyncio.CancelledError):
            # Wait for cancellation
            await t1

        # Make sure only one lock was taken
        self.assertEqual(lock_count, 1)
        # While 3 calls were made to lockit()
        self.assertEqual(call_count, 3)
        self.assertTrue(t1.cancelled() and t2.done())

        # Cleanup the task that is stuck on acquire.
        t3.cancel()
        await asyncio.sleep(0)
        self.assertTrue(t3.cancelled())

    async def test_finished_waiter_cancelled(self):
        lock = asyncio.Lock()

        await lock.acquire()
        self.assertTrue(lock.locked())

        tb = asyncio.create_task(lock.acquire())
        await asyncio.sleep(0)
        self.assertEqual(len(lock._waiters), 1)

        # Create a second waiter, wake up the first, and cancel it.
        # Without the fix, the second was not woken up.
        tc = asyncio.create_task(lock.acquire())
        tb.cancel()
        lock.release()
        await asyncio.sleep(0)

        self.assertTrue(lock.locked())
        self.assertTrue(tb.cancelled())

        # Cleanup
        await tc

    async def test_release_not_acquired(self):
        lock = asyncio.Lock()

        self.assertRaises(RuntimeError, lock.release)

    async def test_release_no_waiters(self):
        lock = asyncio.Lock()
        await lock.acquire()
        self.assertTrue(lock.locked())

        lock.release()
        self.assertFalse(lock.locked())

    async def test_context_manager(self):
        lock = asyncio.Lock()
        self.assertFalse(lock.locked())

        async with lock:
            self.assertTrue(lock.locked())

        self.assertFalse(lock.locked())


class EventTests(unittest.IsolatedAsyncioTestCase):

    def test_repr(self):
        ev = asyncio.Event()
        self.assertTrue(repr(ev).endswith('[unset]>'))
        match = RGX_REPR.match(repr(ev))
        self.assertEqual(match.group('extras'), 'unset')

        ev.set()
        self.assertTrue(repr(ev).endswith('[set]>'))
        self.assertTrue(RGX_REPR.match(repr(ev)))

        ev._waiters.append(mock.Mock())
        self.assertTrue('waiters:1' in repr(ev))
        self.assertTrue(RGX_REPR.match(repr(ev)))

    async def test_wait(self):
        ev = asyncio.Event()
        self.assertFalse(ev.is_set())

        result = []

        async def c1(result):
            if await ev.wait():
                result.append(1)

        async def c2(result):
            if await ev.wait():
                result.append(2)

        async def c3(result):
            if await ev.wait():
                result.append(3)

        t1 = asyncio.create_task(c1(result))
        t2 = asyncio.create_task(c2(result))

        await asyncio.sleep(0)
        self.assertEqual([], result)

        t3 = asyncio.create_task(c3(result))

        ev.set()
        await asyncio.sleep(0)
        self.assertEqual([3, 1, 2], result)

        self.assertTrue(t1.done())
        self.assertIsNone(t1.result())
        self.assertTrue(t2.done())
        self.assertIsNone(t2.result())
        self.assertTrue(t3.done())
        self.assertIsNone(t3.result())

    async def test_wait_on_set(self):
        ev = asyncio.Event()
        ev.set()

        res = await ev.wait()
        self.assertTrue(res)

    async def test_wait_cancel(self):
        ev = asyncio.Event()

        wait = asyncio.create_task(ev.wait())
        asyncio.get_running_loop().call_soon(wait.cancel)
        with self.assertRaises(asyncio.CancelledError):
            await wait
        self.assertFalse(ev._waiters)

    async def test_clear(self):
        ev = asyncio.Event()
        self.assertFalse(ev.is_set())

        ev.set()
        self.assertTrue(ev.is_set())

        ev.clear()
        self.assertFalse(ev.is_set())

    async def test_clear_with_waiters(self):
        ev = asyncio.Event()
        result = []

        async def c1(result):
            if await ev.wait():
                result.append(1)
            return True

        t = asyncio.create_task(c1(result))
        await asyncio.sleep(0)
        self.assertEqual([], result)

        ev.set()
        ev.clear()
        self.assertFalse(ev.is_set())

        ev.set()
        ev.set()
        self.assertEqual(1, len(ev._waiters))

        await asyncio.sleep(0)
        self.assertEqual([1], result)
        self.assertEqual(0, len(ev._waiters))

        self.assertTrue(t.done())
        self.assertTrue(t.result())


class ConditionTests(unittest.IsolatedAsyncioTestCase):

    async def test_wait(self):
        cond = asyncio.Condition()
        result = []

        async def c1(result):
            await cond.acquire()
            if await cond.wait():
                result.append(1)
            return True

        async def c2(result):
            await cond.acquire()
            if await cond.wait():
                result.append(2)
            return True

        async def c3(result):
            await cond.acquire()
            if await cond.wait():
                result.append(3)
            return True

        t1 = asyncio.create_task(c1(result))
        t2 = asyncio.create_task(c2(result))
        t3 = asyncio.create_task(c3(result))

        await asyncio.sleep(0)
        self.assertEqual([], result)
        self.assertFalse(cond.locked())

        self.assertTrue(await cond.acquire())
        cond.notify()
        await asyncio.sleep(0)
        self.assertEqual([], result)
        self.assertTrue(cond.locked())

        cond.release()
        await asyncio.sleep(0)
        self.assertEqual([1], result)
        self.assertTrue(cond.locked())

        cond.notify(2)
        await asyncio.sleep(0)
        self.assertEqual([1], result)
        self.assertTrue(cond.locked())

        cond.release()
        await asyncio.sleep(0)
        self.assertEqual([1, 2], result)
        self.assertTrue(cond.locked())

        cond.release()
        await asyncio.sleep(0)
        self.assertEqual([1, 2, 3], result)
        self.assertTrue(cond.locked())

        self.assertTrue(t1.done())
        self.assertTrue(t1.result())
        self.assertTrue(t2.done())
        self.assertTrue(t2.result())
        self.assertTrue(t3.done())
        self.assertTrue(t3.result())

    async def test_wait_cancel(self):
        cond = asyncio.Condition()
        await cond.acquire()

        wait = asyncio.create_task(cond.wait())
        asyncio.get_running_loop().call_soon(wait.cancel)
        with self.assertRaises(asyncio.CancelledError):
            await wait
        self.assertFalse(cond._waiters)
        self.assertTrue(cond.locked())

    async def test_wait_cancel_contested(self):
        cond = asyncio.Condition()

        await cond.acquire()
        self.assertTrue(cond.locked())

        wait_task = asyncio.create_task(cond.wait())
        await asyncio.sleep(0)
        self.assertFalse(cond.locked())

        # Notify, but contest the lock before cancelling
        await cond.acquire()
        self.assertTrue(cond.locked())
        cond.notify()
        asyncio.get_running_loop().call_soon(wait_task.cancel)
        asyncio.get_running_loop().call_soon(cond.release)

        try:
            await wait_task
        except asyncio.CancelledError:
            # Should not happen, since no cancellation points
            pass

        self.assertTrue(cond.locked())

    async def test_wait_cancel_after_notify(self):
        # See bpo-32841
        waited = False

        cond = asyncio.Condition()

        async def wait_on_cond():
            nonlocal waited
            async with cond:
                waited = True  # Make sure this area was reached
                await cond.wait()

        waiter = asyncio.create_task(wait_on_cond())
        await asyncio.sleep(0)  # Start waiting

        await cond.acquire()
        cond.notify()
        await asyncio.sleep(0)  # Get to acquire()
        waiter.cancel()
        await asyncio.sleep(0)  # Activate cancellation
        cond.release()
        await asyncio.sleep(0)  # Cancellation should occur

        self.assertTrue(waiter.cancelled())
        self.assertTrue(waited)

    async def test_wait_unacquired(self):
        cond = asyncio.Condition()
        with self.assertRaises(RuntimeError):
            await cond.wait()

    async def test_wait_for(self):
        cond = asyncio.Condition()
        presult = False

        def predicate():
            return presult

        result = []

        async def c1(result):
            await cond.acquire()
            if await cond.wait_for(predicate):
                result.append(1)
                cond.release()
            return True

        t = asyncio.create_task(c1(result))

        await asyncio.sleep(0)
        self.assertEqual([], result)

        await cond.acquire()
        cond.notify()
        cond.release()
        await asyncio.sleep(0)
        self.assertEqual([], result)

        presult = True
        await cond.acquire()
        cond.notify()
        cond.release()
        await asyncio.sleep(0)
        self.assertEqual([1], result)

        self.assertTrue(t.done())
        self.assertTrue(t.result())

    async def test_wait_for_unacquired(self):
        cond = asyncio.Condition()

        # predicate can return true immediately
        res = await cond.wait_for(lambda: [1, 2, 3])
        self.assertEqual([1, 2, 3], res)

        with self.assertRaises(RuntimeError):
            await cond.wait_for(lambda: False)

    async def test_notify(self):
        cond = asyncio.Condition()
        result = []

        async def c1(result):
            await cond.acquire()
            if await cond.wait():
                result.append(1)
                cond.release()
            return True

        async def c2(result):
            await cond.acquire()
            if await cond.wait():
                result.append(2)
                cond.release()
            return True

        async def c3(result):
            await cond.acquire()
            if await cond.wait():
                result.append(3)
                cond.release()
            return True

        t1 = asyncio.create_task(c1(result))
        t2 = asyncio.create_task(c2(result))
        t3 = asyncio.create_task(c3(result))

        await asyncio.sleep(0)
        self.assertEqual([], result)

        await cond.acquire()
        cond.notify(1)
        cond.release()
        await asyncio.sleep(0)
        self.assertEqual([1], result)

        await cond.acquire()
        cond.notify(1)
        cond.notify(2048)
        cond.release()
        await asyncio.sleep(0)
        self.assertEqual([1, 2, 3], result)

        self.assertTrue(t1.done())
        self.assertTrue(t1.result())
        self.assertTrue(t2.done())
        self.assertTrue(t2.result())
        self.assertTrue(t3.done())
        self.assertTrue(t3.result())

    async def test_notify_all(self):
        cond = asyncio.Condition()

        result = []

        async def c1(result):
            await cond.acquire()
            if await cond.wait():
                result.append(1)
                cond.release()
            return True

        async def c2(result):
            await cond.acquire()
            if await cond.wait():
                result.append(2)
                cond.release()
            return True

        t1 = asyncio.create_task(c1(result))
        t2 = asyncio.create_task(c2(result))

        await asyncio.sleep(0)
        self.assertEqual([], result)

        await cond.acquire()
        cond.notify_all()
        cond.release()
        await asyncio.sleep(0)
        self.assertEqual([1, 2], result)

        self.assertTrue(t1.done())
        self.assertTrue(t1.result())
        self.assertTrue(t2.done())
        self.assertTrue(t2.result())

    def test_notify_unacquired(self):
        cond = asyncio.Condition()
        self.assertRaises(RuntimeError, cond.notify)

    def test_notify_all_unacquired(self):
        cond = asyncio.Condition()
        self.assertRaises(RuntimeError, cond.notify_all)

    async def test_repr(self):
        cond = asyncio.Condition()
        self.assertTrue('unlocked' in repr(cond))
        self.assertTrue(RGX_REPR.match(repr(cond)))

        await cond.acquire()
        self.assertTrue('locked' in repr(cond))

        cond._waiters.append(mock.Mock())
        self.assertTrue('waiters:1' in repr(cond))
        self.assertTrue(RGX_REPR.match(repr(cond)))

        cond._waiters.append(mock.Mock())
        self.assertTrue('waiters:2' in repr(cond))
        self.assertTrue(RGX_REPR.match(repr(cond)))

    async def test_context_manager(self):
        cond = asyncio.Condition()
        self.assertFalse(cond.locked())
        async with cond:
            self.assertTrue(cond.locked())
        self.assertFalse(cond.locked())

    async def test_explicit_lock(self):
        async def f(lock=None, cond=None):
            if lock is None:
                lock = asyncio.Lock()
            if cond is None:
                cond = asyncio.Condition(lock)
            self.assertIs(cond._lock, lock)
            self.assertFalse(lock.locked())
            self.assertFalse(cond.locked())
            async with cond:
                self.assertTrue(lock.locked())
                self.assertTrue(cond.locked())
            self.assertFalse(lock.locked())
            self.assertFalse(cond.locked())
            async with lock:
                self.assertTrue(lock.locked())
                self.assertTrue(cond.locked())
            self.assertFalse(lock.locked())
            self.assertFalse(cond.locked())

        # All should work in the same way.
        await f()
        await f(asyncio.Lock())
        lock = asyncio.Lock()
        await f(lock, asyncio.Condition(lock))

    async def test_ambiguous_loops(self):
        loop = asyncio.new_event_loop()
        self.addCleanup(loop.close)

        async def wrong_loop_in_lock():
            with self.assertRaises(TypeError):
                asyncio.Lock(loop=loop)  # actively disallowed since 3.10
            lock = asyncio.Lock()
            lock._loop = loop  # use private API for testing
            async with lock:
                # acquired immediately via the fast-path
                # without interaction with any event loop.
                cond = asyncio.Condition(lock)
                # cond.acquire() will trigger waiting on the lock
                # and it will discover the event loop mismatch.
                with self.assertRaisesRegex(
                    RuntimeError,
                    "is bound to a different event loop",
                ):
                    await cond.acquire()

        async def wrong_loop_in_cond():
            # Same analogy here with the condition's loop.
            lock = asyncio.Lock()
            async with lock:
                with self.assertRaises(TypeError):
                    asyncio.Condition(lock, loop=loop)
                cond = asyncio.Condition(lock)
                cond._loop = loop
                with self.assertRaisesRegex(
                    RuntimeError,
                    "is bound to a different event loop",
                ):
                    await cond.wait()

        await wrong_loop_in_lock()
        await wrong_loop_in_cond()

    async def test_timeout_in_block(self):
        condition = asyncio.Condition()
        async with condition:
            with self.assertRaises(asyncio.TimeoutError):
                await asyncio.wait_for(condition.wait(), timeout=0.5)

    async def test_cancelled_error_wakeup(self):
        # Test that a cancelled error, received when awaiting wakeup,
        # will be re-raised un-modified.
        wake = False
        raised = None
        cond = asyncio.Condition()

        async def func():
            nonlocal raised
            async with cond:
                with self.assertRaises(asyncio.CancelledError) as err:
                    await cond.wait_for(lambda: wake)
                raised = err.exception
                raise raised

        task = asyncio.create_task(func())
        await asyncio.sleep(0)
        # Task is waiting on the condition, cancel it there.
        task.cancel(msg="foo")
        with self.assertRaises(asyncio.CancelledError) as err:
            await task
        self.assertEqual(err.exception.args, ("foo",))
        # We should have got the _same_ exception instance as the one
        # originally raised.
        self.assertIs(err.exception, raised)

    async def test_cancelled_error_re_aquire(self):
        # Test that a cancelled error, received when re-aquiring lock,
        # will be re-raised un-modified.
        wake = False
        raised = None
        cond = asyncio.Condition()

        async def func():
            nonlocal raised
            async with cond:
                with self.assertRaises(asyncio.CancelledError) as err:
                    await cond.wait_for(lambda: wake)
                raised = err.exception
                raise raised

        task = asyncio.create_task(func())
        await asyncio.sleep(0)
        # Task is waiting on the condition
        await cond.acquire()
        wake = True
        cond.notify()
        await asyncio.sleep(0)
        # Task is now trying to re-acquire the lock, cancel it there.
        task.cancel(msg="foo")
        cond.release()
        with self.assertRaises(asyncio.CancelledError) as err:
            await task
        self.assertEqual(err.exception.args, ("foo",))
        # We should have got the _same_ exception instance as the one
        # originally raised.
        self.assertIs(err.exception, raised)

    async def test_cancelled_wakeup(self):
        # Test that a task cancelled at the "same" time as it is woken
        # up as part of a Condition.notify() does not result in a lost wakeup.
        # This test simulates a cancel while the target task is awaiting initial
        # wakeup on the wakeup queue.
        condition = asyncio.Condition()
        state = 0
        async def consumer():
            nonlocal state
            async with condition:
                while True:
                    await condition.wait_for(lambda: state != 0)
                    if state < 0:
                        return
                    state -= 1

        # create two consumers
        c = [asyncio.create_task(consumer()) for _ in range(2)]
        # wait for them to settle
        await asyncio.sleep(0)
        async with condition:
            # produce one item and wake up one
            state += 1
            condition.notify(1)

            # Cancel it while it is awaiting to be run.
            # This cancellation could come from the outside
            c[0].cancel()

            # now wait for the item to be consumed
            # if it doesn't means that our "notify" didn"t take hold.
            # because it raced with a cancel()
            try:
                async with asyncio.timeout(0.01):
                    await condition.wait_for(lambda: state == 0)
            except TimeoutError:
                pass
            self.assertEqual(state, 0)

            # clean up
            state = -1
            condition.notify_all()
        await c[1]

    async def test_cancelled_wakeup_relock(self):
        # Test that a task cancelled at the "same" time as it is woken
        # up as part of a Condition.notify() does not result in a lost wakeup.
        # This test simulates a cancel while the target task is acquiring the lock
        # again.
        condition = asyncio.Condition()
        state = 0
        async def consumer():
            nonlocal state
            async with condition:
                while True:
                    await condition.wait_for(lambda: state != 0)
                    if state < 0:
                        return
                    state -= 1

        # create two consumers
        c = [asyncio.create_task(consumer()) for _ in range(2)]
        # wait for them to settle
        await asyncio.sleep(0)
        async with condition:
            # produce one item and wake up one
            state += 1
            condition.notify(1)

            # now we sleep for a bit.  This allows the target task to wake up and
            # settle on re-aquiring the lock
            await asyncio.sleep(0)

            # Cancel it while awaiting the lock
            # This cancel could come the outside.
            c[0].cancel()

            # now wait for the item to be consumed
            # if it doesn't means that our "notify" didn"t take hold.
            # because it raced with a cancel()
            try:
                async with asyncio.timeout(0.01):
                    await condition.wait_for(lambda: state == 0)
            except TimeoutError:
                pass
            self.assertEqual(state, 0)

            # clean up
            state = -1
            condition.notify_all()
        await c[1]

class SemaphoreTests(unittest.IsolatedAsyncioTestCase):

    def test_initial_value_zero(self):
        sem = asyncio.Semaphore(0)
        self.assertTrue(sem.locked())

    async def test_repr(self):
        sem = asyncio.Semaphore()
        self.assertTrue(repr(sem).endswith('[unlocked, value:1]>'))
        self.assertTrue(RGX_REPR.match(repr(sem)))

        await sem.acquire()
        self.assertTrue(repr(sem).endswith('[locked]>'))
        self.assertTrue('waiters' not in repr(sem))
        self.assertTrue(RGX_REPR.match(repr(sem)))

        if sem._waiters is None:
            sem._waiters = collections.deque()

        sem._waiters.append(mock.Mock())
        self.assertTrue('waiters:1' in repr(sem))
        self.assertTrue(RGX_REPR.match(repr(sem)))

        sem._waiters.append(mock.Mock())
        self.assertTrue('waiters:2' in repr(sem))
        self.assertTrue(RGX_REPR.match(repr(sem)))

    async def test_semaphore(self):
        sem = asyncio.Semaphore()
        self.assertEqual(1, sem._value)

        with self.assertRaisesRegex(
            TypeError,
            "'Semaphore' object can't be awaited",
        ):
            await sem

        self.assertFalse(sem.locked())
        self.assertEqual(1, sem._value)

    def test_semaphore_value(self):
        self.assertRaises(ValueError, asyncio.Semaphore, -1)

    async def test_acquire(self):
        sem = asyncio.Semaphore(3)
        result = []

        self.assertTrue(await sem.acquire())
        self.assertTrue(await sem.acquire())
        self.assertFalse(sem.locked())

        async def c1(result):
            await sem.acquire()
            result.append(1)
            return True

        async def c2(result):
            await sem.acquire()
            result.append(2)
            return True

        async def c3(result):
            await sem.acquire()
            result.append(3)
            return True

        async def c4(result):
            await sem.acquire()
            result.append(4)
            return True

        t1 = asyncio.create_task(c1(result))
        t2 = asyncio.create_task(c2(result))
        t3 = asyncio.create_task(c3(result))

        await asyncio.sleep(0)
        self.assertEqual([1], result)
        self.assertTrue(sem.locked())
        self.assertEqual(2, len(sem._waiters))
        self.assertEqual(0, sem._value)

        t4 = asyncio.create_task(c4(result))

        sem.release()
        sem.release()
        self.assertEqual(0, sem._value)

        await asyncio.sleep(0)
        self.assertEqual(0, sem._value)
        self.assertEqual(3, len(result))
        self.assertTrue(sem.locked())
        self.assertEqual(1, len(sem._waiters))
        self.assertEqual(0, sem._value)

        self.assertTrue(t1.done())
        self.assertTrue(t1.result())
        race_tasks = [t2, t3, t4]
        done_tasks = [t for t in race_tasks if t.done() and t.result()]
        self.assertEqual(2, len(done_tasks))

        # cleanup locked semaphore
        sem.release()
        await asyncio.gather(*race_tasks)

    async def test_acquire_cancel(self):
        sem = asyncio.Semaphore()
        await sem.acquire()

        acquire = asyncio.create_task(sem.acquire())
        asyncio.get_running_loop().call_soon(acquire.cancel)
        with self.assertRaises(asyncio.CancelledError):
            await acquire
        self.assertTrue((not sem._waiters) or
                        all(waiter.done() for waiter in sem._waiters))

    async def test_acquire_cancel_before_awoken(self):
        sem = asyncio.Semaphore(value=0)

        t1 = asyncio.create_task(sem.acquire())
        t2 = asyncio.create_task(sem.acquire())
        t3 = asyncio.create_task(sem.acquire())
        t4 = asyncio.create_task(sem.acquire())

        await asyncio.sleep(0)

        t1.cancel()
        t2.cancel()
        sem.release()

        await asyncio.sleep(0)
        await asyncio.sleep(0)
        num_done = sum(t.done() for t in [t3, t4])
        self.assertEqual(num_done, 1)
        self.assertTrue(t3.done())
        self.assertFalse(t4.done())

        t3.cancel()
        t4.cancel()
        await asyncio.sleep(0)

    async def test_acquire_hang(self):
        sem = asyncio.Semaphore(value=0)

        t1 = asyncio.create_task(sem.acquire())
        t2 = asyncio.create_task(sem.acquire())
        await asyncio.sleep(0)

        t1.cancel()
        sem.release()
        await asyncio.sleep(0)
        await asyncio.sleep(0)
        self.assertTrue(sem.locked())
        self.assertTrue(t2.done())

    async def test_acquire_no_hang(self):

        sem = asyncio.Semaphore(1)

        async def c1():
            async with sem:
                await asyncio.sleep(0)
            t2.cancel()

        async def c2():
            async with sem:
                self.assertFalse(True)

        t1 = asyncio.create_task(c1())
        t2 = asyncio.create_task(c2())

        r1, r2 = await asyncio.gather(t1, t2, return_exceptions=True)
        self.assertTrue(r1 is None)
        self.assertTrue(isinstance(r2, asyncio.CancelledError))

        await asyncio.wait_for(sem.acquire(), timeout=1.0)

    def test_release_not_acquired(self):
        sem = asyncio.BoundedSemaphore()

        self.assertRaises(ValueError, sem.release)

    async def test_release_no_waiters(self):
        sem = asyncio.Semaphore()
        await sem.acquire()
        self.assertTrue(sem.locked())

        sem.release()
        self.assertFalse(sem.locked())

    async def test_acquire_fifo_order(self):
        sem = asyncio.Semaphore(1)
        result = []

        async def coro(tag):
            await sem.acquire()
            result.append(f'{tag}_1')
            await asyncio.sleep(0.01)
            sem.release()

            await sem.acquire()
            result.append(f'{tag}_2')
            await asyncio.sleep(0.01)
            sem.release()

        async with asyncio.TaskGroup() as tg:
            tg.create_task(coro('c1'))
            tg.create_task(coro('c2'))
            tg.create_task(coro('c3'))

        self.assertEqual(
            ['c1_1', 'c2_1', 'c3_1', 'c1_2', 'c2_2', 'c3_2'],
            result
        )

    async def test_acquire_fifo_order_2(self):
        sem = asyncio.Semaphore(1)
        result = []

        async def c1(result):
            await sem.acquire()
            result.append(1)
            return True

        async def c2(result):
            await sem.acquire()
            result.append(2)
            sem.release()
            await sem.acquire()
            result.append(4)
            return True

        async def c3(result):
            await sem.acquire()
            result.append(3)
            return True

        t1 = asyncio.create_task(c1(result))
        t2 = asyncio.create_task(c2(result))
        t3 = asyncio.create_task(c3(result))

        await asyncio.sleep(0)

        sem.release()
        sem.release()

        tasks = [t1, t2, t3]
        await asyncio.gather(*tasks)
        self.assertEqual([1, 2, 3, 4], result)

    async def test_acquire_fifo_order_3(self):
        sem = asyncio.Semaphore(0)
        result = []

        async def c1(result):
            await sem.acquire()
            result.append(1)
            return True

        async def c2(result):
            await sem.acquire()
            result.append(2)
            return True

        async def c3(result):
            await sem.acquire()
            result.append(3)
            return True

        t1 = asyncio.create_task(c1(result))
        t2 = asyncio.create_task(c2(result))
        t3 = asyncio.create_task(c3(result))

        await asyncio.sleep(0)

        t1.cancel()

        await asyncio.sleep(0)

        sem.release()
        sem.release()

        tasks = [t1, t2, t3]
        await asyncio.gather(*tasks, return_exceptions=True)
        self.assertEqual([2, 3], result)

    async def test_acquire_fifo_order_4(self):
        # Test that a successfule `acquire()` will wake up multiple Tasks
        # that were waiting in the Semaphore queue due to FIFO rules.
        sem = asyncio.Semaphore(0)
        result = []
        count = 0

        async def c1(result):
            # First task immediatlly waits for semaphore.  It will be awoken by c2.
            self.assertEqual(sem._value, 0)
            await sem.acquire()
            # We should have woken up all waiting tasks now.
            self.assertEqual(sem._value, 0)
            # Create a fourth task.  It should run after c3, not c2.
            nonlocal t4
            t4 = asyncio.create_task(c4(result))
            result.append(1)
            return True

        async def c2(result):
            # The second task begins by releasing semaphore three times,
            # for c1, c2, and c3.
            sem.release()
            sem.release()
            sem.release()
            self.assertEqual(sem._value, 2)
            # It is locked, because c1 hasn't woken up yet.
            self.assertTrue(sem.locked())
            await sem.acquire()
            result.append(2)
            return True

        async def c3(result):
            await sem.acquire()
            self.assertTrue(sem.locked())
            result.append(3)
            return True

        async def c4(result):
            result.append(4)
            return True

        t1 = asyncio.create_task(c1(result))
        t2 = asyncio.create_task(c2(result))
        t3 = asyncio.create_task(c3(result))
        t4 = None

        await asyncio.sleep(0)
        # Three tasks are in the queue, the first hasn't woken up yet.
        self.assertEqual(sem._value, 2)
        self.assertEqual(len(sem._waiters), 3)
        await asyncio.sleep(0)

        tasks = [t1, t2, t3, t4]
        await asyncio.gather(*tasks)
        self.assertEqual([1, 2, 3, 4], result)

class BarrierTests(unittest.IsolatedAsyncioTestCase):

    async def asyncSetUp(self):
        await super().asyncSetUp()
        self.N = 5

    def make_tasks(self, n, coro):
        tasks = [asyncio.create_task(coro()) for _ in range(n)]
        return tasks

    async def gather_tasks(self, n, coro):
        tasks = self.make_tasks(n, coro)
        res = await asyncio.gather(*tasks)
        return res, tasks

    async def test_barrier(self):
        barrier = asyncio.Barrier(self.N)
        self.assertIn("filling", repr(barrier))
        with self.assertRaisesRegex(
            TypeError,
            "'Barrier' object can't be awaited",
        ):
            await barrier

        self.assertIn("filling", repr(barrier))

    async def test_repr(self):
        barrier = asyncio.Barrier(self.N)

        self.assertTrue(RGX_REPR.match(repr(barrier)))
        self.assertIn("filling", repr(barrier))

        waiters = []
        async def wait(barrier):
            await barrier.wait()

        incr = 2
        for i in range(incr):
            waiters.append(asyncio.create_task(wait(barrier)))
        await asyncio.sleep(0)

        self.assertTrue(RGX_REPR.match(repr(barrier)))
        self.assertTrue(f"waiters:{incr}/{self.N}" in repr(barrier))
        self.assertIn("filling", repr(barrier))

        # create missing waiters
        for i in range(barrier.parties - barrier.n_waiting):
            waiters.append(asyncio.create_task(wait(barrier)))
        await asyncio.sleep(0)

        self.assertTrue(RGX_REPR.match(repr(barrier)))
        self.assertIn("draining", repr(barrier))

        # add a part of waiters
        for i in range(incr):
            waiters.append(asyncio.create_task(wait(barrier)))
        await asyncio.sleep(0)
        # and reset
        await barrier.reset()

        self.assertTrue(RGX_REPR.match(repr(barrier)))
        self.assertIn("resetting", repr(barrier))

        # add a part of waiters again
        for i in range(incr):
            waiters.append(asyncio.create_task(wait(barrier)))
        await asyncio.sleep(0)
        # and abort
        await barrier.abort()

        self.assertTrue(RGX_REPR.match(repr(barrier)))
        self.assertIn("broken", repr(barrier))
        self.assertTrue(barrier.broken)

        # suppress unhandled exceptions
        await asyncio.gather(*waiters, return_exceptions=True)

    async def test_barrier_parties(self):
        self.assertRaises(ValueError, lambda: asyncio.Barrier(0))
        self.assertRaises(ValueError, lambda: asyncio.Barrier(-4))

        self.assertIsInstance(asyncio.Barrier(self.N), asyncio.Barrier)

    async def test_context_manager(self):
        self.N = 3
        barrier = asyncio.Barrier(self.N)
        results = []

        async def coro():
            async with barrier as i:
                results.append(i)

        await self.gather_tasks(self.N, coro)

        self.assertListEqual(sorted(results), list(range(self.N)))
        self.assertEqual(barrier.n_waiting, 0)
        self.assertFalse(barrier.broken)

    async def test_filling_one_task(self):
        barrier = asyncio.Barrier(1)

        async def f():
            async with barrier as i:
                return True

        ret = await f()

        self.assertTrue(ret)
        self.assertEqual(barrier.n_waiting, 0)
        self.assertFalse(barrier.broken)

    async def test_filling_one_task_twice(self):
        barrier = asyncio.Barrier(1)

        t1 = asyncio.create_task(barrier.wait())
        await asyncio.sleep(0)
        self.assertEqual(barrier.n_waiting, 0)

        t2 = asyncio.create_task(barrier.wait())
        await asyncio.sleep(0)

        self.assertEqual(t1.result(), t2.result())
        self.assertEqual(t1.done(), t2.done())

        self.assertEqual(barrier.n_waiting, 0)
        self.assertFalse(barrier.broken)

    async def test_filling_task_by_task(self):
        self.N = 3
        barrier = asyncio.Barrier(self.N)

        t1 = asyncio.create_task(barrier.wait())
        await asyncio.sleep(0)
        self.assertEqual(barrier.n_waiting, 1)
        self.assertIn("filling", repr(barrier))

        t2 = asyncio.create_task(barrier.wait())
        await asyncio.sleep(0)
        self.assertEqual(barrier.n_waiting, 2)
        self.assertIn("filling", repr(barrier))

        t3 = asyncio.create_task(barrier.wait())
        await asyncio.sleep(0)

        await asyncio.wait([t1, t2, t3])

        self.assertEqual(barrier.n_waiting, 0)
        self.assertFalse(barrier.broken)

    async def test_filling_tasks_wait_twice(self):
        barrier = asyncio.Barrier(self.N)
        results = []

        async def coro():
            async with barrier:
                results.append(True)

                async with barrier:
                    results.append(False)

        await self.gather_tasks(self.N, coro)

        self.assertEqual(len(results), self.N*2)
        self.assertEqual(results.count(True), self.N)
        self.assertEqual(results.count(False), self.N)

        self.assertEqual(barrier.n_waiting, 0)
        self.assertFalse(barrier.broken)

    async def test_filling_tasks_check_return_value(self):
        barrier = asyncio.Barrier(self.N)
        results1 = []
        results2 = []

        async def coro():
            async with barrier:
                results1.append(True)

                async with barrier as i:
                    results2.append(True)
                    return i

        res, _ = await self.gather_tasks(self.N, coro)

        self.assertEqual(len(results1), self.N)
        self.assertTrue(all(results1))
        self.assertEqual(len(results2), self.N)
        self.assertTrue(all(results2))
        self.assertListEqual(sorted(res), list(range(self.N)))

        self.assertEqual(barrier.n_waiting, 0)
        self.assertFalse(barrier.broken)

    async def test_draining_state(self):
        barrier = asyncio.Barrier(self.N)
        results = []

        async def coro():
            async with barrier:
                # barrier state change to filling for the last task release
                results.append("draining" in repr(barrier))

        await self.gather_tasks(self.N, coro)

        self.assertEqual(len(results), self.N)
        self.assertEqual(results[-1], False)
        self.assertTrue(all(results[:self.N-1]))

        self.assertEqual(barrier.n_waiting, 0)
        self.assertFalse(barrier.broken)

    async def test_blocking_tasks_while_draining(self):
        rewait = 2
        barrier = asyncio.Barrier(self.N)
        barrier_nowaiting = asyncio.Barrier(self.N - rewait)
        results = []
        rewait_n = rewait
        counter = 0

        async def coro():
            nonlocal rewait_n

            # first time waiting
            await barrier.wait()

            # after wainting once for all tasks
            if rewait_n > 0:
                rewait_n -= 1
                # wait again only for rewait tasks
                await barrier.wait()
            else:
                # wait for end of draining state
                await barrier_nowaiting.wait()
                # wait for other waiting tasks
                await barrier.wait()

        # a success means that barrier_nowaiting
        # was waited for exactly N-rewait=3 times
        await self.gather_tasks(self.N, coro)

    async def test_filling_tasks_cancel_one(self):
        self.N = 3
        barrier = asyncio.Barrier(self.N)
        results = []

        async def coro():
            await barrier.wait()
            results.append(True)

        t1 = asyncio.create_task(coro())
        await asyncio.sleep(0)
        self.assertEqual(barrier.n_waiting, 1)

        t2 = asyncio.create_task(coro())
        await asyncio.sleep(0)
        self.assertEqual(barrier.n_waiting, 2)

        t1.cancel()
        await asyncio.sleep(0)
        self.assertEqual(barrier.n_waiting, 1)
        with self.assertRaises(asyncio.CancelledError):
            await t1
        self.assertTrue(t1.cancelled())

        t3 = asyncio.create_task(coro())
        await asyncio.sleep(0)
        self.assertEqual(barrier.n_waiting, 2)

        t4 = asyncio.create_task(coro())
        await asyncio.gather(t2, t3, t4)

        self.assertEqual(len(results), self.N)
        self.assertTrue(all(results))

        self.assertEqual(barrier.n_waiting, 0)
        self.assertFalse(barrier.broken)

    async def test_reset_barrier(self):
        barrier = asyncio.Barrier(1)

        asyncio.create_task(barrier.reset())
        await asyncio.sleep(0)

        self.assertEqual(barrier.n_waiting, 0)
        self.assertFalse(barrier.broken)

    async def test_reset_barrier_while_tasks_waiting(self):
        barrier = asyncio.Barrier(self.N)
        results = []

        async def coro():
            try:
                await barrier.wait()
            except asyncio.BrokenBarrierError:
                results.append(True)

        async def coro_reset():
            await barrier.reset()

        # N-1 tasks waiting on barrier with N parties
        tasks  = self.make_tasks(self.N-1, coro)
        await asyncio.sleep(0)

        # reset the barrier
        asyncio.create_task(coro_reset())
        await asyncio.gather(*tasks)

        self.assertEqual(len(results), self.N-1)
        self.assertTrue(all(results))
        self.assertEqual(barrier.n_waiting, 0)
        self.assertNotIn("resetting", repr(barrier))
        self.assertFalse(barrier.broken)

    async def test_reset_barrier_when_tasks_half_draining(self):
        barrier = asyncio.Barrier(self.N)
        results1 = []
        rest_of_tasks = self.N//2

        async def coro():
            try:
                await barrier.wait()
            except asyncio.BrokenBarrierError:
                # catch here waiting tasks
                results1.append(True)
            else:
                # here drained task outside the barrier
                if rest_of_tasks == barrier._count:
                    # tasks outside the barrier
                    await barrier.reset()

        await self.gather_tasks(self.N, coro)

        self.assertEqual(results1, [True]*rest_of_tasks)
        self.assertEqual(barrier.n_waiting, 0)
        self.assertNotIn("resetting", repr(barrier))
        self.assertFalse(barrier.broken)

    async def test_reset_barrier_when_tasks_half_draining_half_blocking(self):
        barrier = asyncio.Barrier(self.N)
        results1 = []
        results2 = []
        blocking_tasks = self.N//2
        count = 0

        async def coro():
            nonlocal count
            try:
                await barrier.wait()
            except asyncio.BrokenBarrierError:
                # here catch still waiting tasks
                results1.append(True)

                # so now waiting again to reach nb_parties
                await barrier.wait()
            else:
                count += 1
                if count > blocking_tasks:
                    # reset now: raise asyncio.BrokenBarrierError for waiting tasks
                    await barrier.reset()

                    # so now waiting again to reach nb_parties
                    await barrier.wait()
                else:
                    try:
                        await barrier.wait()
                    except asyncio.BrokenBarrierError:
                        # here no catch - blocked tasks go to wait
                        results2.append(True)

        await self.gather_tasks(self.N, coro)

        self.assertEqual(results1, [True]*blocking_tasks)
        self.assertEqual(results2, [])
        self.assertEqual(barrier.n_waiting, 0)
        self.assertNotIn("resetting", repr(barrier))
        self.assertFalse(barrier.broken)

    async def test_reset_barrier_while_tasks_waiting_and_waiting_again(self):
        barrier = asyncio.Barrier(self.N)
        results1 = []
        results2 = []

        async def coro1():
            try:
                await barrier.wait()
            except asyncio.BrokenBarrierError:
                results1.append(True)
            finally:
                await barrier.wait()
                results2.append(True)

        async def coro2():
            async with barrier:
                results2.append(True)

        tasks = self.make_tasks(self.N-1, coro1)

        # reset barrier, N-1 waiting tasks raise an BrokenBarrierError
        asyncio.create_task(barrier.reset())
        await asyncio.sleep(0)

        # complete waiting tasks in the `finally`
        asyncio.create_task(coro2())

        await asyncio.gather(*tasks)

        self.assertFalse(barrier.broken)
        self.assertEqual(len(results1), self.N-1)
        self.assertTrue(all(results1))
        self.assertEqual(len(results2), self.N)
        self.assertTrue(all(results2))

        self.assertEqual(barrier.n_waiting, 0)


    async def test_reset_barrier_while_tasks_draining(self):
        barrier = asyncio.Barrier(self.N)
        results1 = []
        results2 = []
        results3 = []
        count = 0

        async def coro():
            nonlocal count

            i = await barrier.wait()
            count += 1
            if count == self.N:
                # last task exited from barrier
                await barrier.reset()

                # wait here to reach the `parties`
                await barrier.wait()
            else:
                try:
                    # second waiting
                    await barrier.wait()

                    # N-1 tasks here
                    results1.append(True)
                except Exception as e:
                    # never goes here
                    results2.append(True)

            # Now, pass the barrier again
            # last wait, must be completed
            k = await barrier.wait()
            results3.append(True)

        await self.gather_tasks(self.N, coro)

        self.assertFalse(barrier.broken)
        self.assertTrue(all(results1))
        self.assertEqual(len(results1), self.N-1)
        self.assertEqual(len(results2), 0)
        self.assertEqual(len(results3), self.N)
        self.assertTrue(all(results3))

        self.assertEqual(barrier.n_waiting, 0)

    async def test_abort_barrier(self):
        barrier = asyncio.Barrier(1)

        asyncio.create_task(barrier.abort())
        await asyncio.sleep(0)

        self.assertEqual(barrier.n_waiting, 0)
        self.assertTrue(barrier.broken)

    async def test_abort_barrier_when_tasks_half_draining_half_blocking(self):
        barrier = asyncio.Barrier(self.N)
        results1 = []
        results2 = []
        blocking_tasks = self.N//2
        count = 0

        async def coro():
            nonlocal count
            try:
                await barrier.wait()
            except asyncio.BrokenBarrierError:
                # here catch tasks waiting to drain
                results1.append(True)
            else:
                count += 1
                if count > blocking_tasks:
                    # abort now: raise asyncio.BrokenBarrierError for all tasks
                    await barrier.abort()
                else:
                    try:
                        await barrier.wait()
                    except asyncio.BrokenBarrierError:
                        # here catch blocked tasks (already drained)
                        results2.append(True)

        await self.gather_tasks(self.N, coro)

        self.assertTrue(barrier.broken)
        self.assertEqual(results1, [True]*blocking_tasks)
        self.assertEqual(results2, [True]*(self.N-blocking_tasks-1))
        self.assertEqual(barrier.n_waiting, 0)
        self.assertNotIn("resetting", repr(barrier))

    async def test_abort_barrier_when_exception(self):
        # test from threading.Barrier: see `lock_tests.test_reset`
        barrier = asyncio.Barrier(self.N)
        results1 = []
        results2 = []

        async def coro():
            try:
                async with barrier as i :
                    if i == self.N//2:
                        raise RuntimeError
                async with barrier:
                    results1.append(True)
            except asyncio.BrokenBarrierError:
                results2.append(True)
            except RuntimeError:
                await barrier.abort()

        await self.gather_tasks(self.N, coro)

        self.assertTrue(barrier.broken)
        self.assertEqual(len(results1), 0)
        self.assertEqual(len(results2), self.N-1)
        self.assertTrue(all(results2))
        self.assertEqual(barrier.n_waiting, 0)

    async def test_abort_barrier_when_exception_then_resetting(self):
        # test from threading.Barrier: see `lock_tests.test_abort_and_reset`
        barrier1 = asyncio.Barrier(self.N)
        barrier2 = asyncio.Barrier(self.N)
        results1 = []
        results2 = []
        results3 = []

        async def coro():
            try:
                i = await barrier1.wait()
                if i == self.N//2:
                    raise RuntimeError
                await barrier1.wait()
                results1.append(True)
            except asyncio.BrokenBarrierError:
                results2.append(True)
            except RuntimeError:
                await barrier1.abort()

            # Synchronize and reset the barrier.  Must synchronize first so
            # that everyone has left it when we reset, and after so that no
            # one enters it before the reset.
            i = await barrier2.wait()
            if  i == self.N//2:
                await barrier1.reset()
            await barrier2.wait()
            await barrier1.wait()
            results3.append(True)

        await self.gather_tasks(self.N, coro)

        self.assertFalse(barrier1.broken)
        self.assertEqual(len(results1), 0)
        self.assertEqual(len(results2), self.N-1)
        self.assertTrue(all(results2))
        self.assertEqual(len(results3), self.N)
        self.assertTrue(all(results3))

        self.assertEqual(barrier1.n_waiting, 0)


if __name__ == '__main__':
    unittest.main()