"""Unit tests for run_limited.

Covers the parts callers depend on: results come back in input order, no more
than `limit` awaitables run at once, and neither a failure nor a cancellation of
the caller leaves anything running in the background.
"""
import asyncio
import unittest

from preseem.util import run_limited


class TestRunLimited(unittest.TestCase):

    def setUp(self):
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)

    def tearDown(self):
        self.loop.close()

    def _await(self, co):
        return self.loop.run_until_complete(co)

    def test_empty_input(self):
        self.assertEqual(self._await(run_limited([], 4)), [])

    def test_results_in_input_order(self):

        async def value(n, delay):
            await asyncio.sleep(delay)
            return n

        # later items finish first, results must still be in input order
        coros = [value(n, delay) for n, delay in enumerate([0.03, 0.02, 0.01, 0])]
        self.assertEqual(self._await(run_limited(coros, 4)), [0, 1, 2, 3])

    def test_respects_concurrency_limit(self):
        state = {'running': 0, 'peak': 0}

        async def track():
            state['running'] += 1
            state['peak'] = max(state['peak'], state['running'])
            await asyncio.sleep(0.01)
            state['running'] -= 1

        self._await(run_limited([track() for _ in range(20)], 3))
        self.assertEqual(state['peak'], 3)
        self.assertEqual(state['running'], 0)

    def test_rejects_nonpositive_limit(self):
        # limit of 0 would deadlock, so it is rejected up front
        for limit in (0, -1):
            self.assertRaises(ValueError, self._await, run_limited([], limit))

    def test_nonpositive_limit_closes_coroutines(self):

        async def never():
            pass

        coro = never()
        self.assertRaises(ValueError, self._await, run_limited([coro], 0))
        # closed, so the ValueError is not buried under never-awaited warnings
        self.assertIsNone(coro.cr_frame)

    def test_first_error_propagates(self):

        async def boom():
            raise ValueError('boom')

        self.assertRaises(ValueError, self._await, run_limited([boom()], 4))

    def test_error_cancels_running_siblings(self):
        unwound = []

        async def boom():
            await asyncio.sleep(0.01)
            raise ValueError('boom')

        async def slow(n):
            try:
                await asyncio.sleep(30)
            except asyncio.CancelledError:
                unwound.append(n)
                raise

        self.assertRaises(ValueError, self._await,
                          run_limited([slow(1), boom(), slow(2)], 4))
        # siblings are cancelled *and* fully unwound before we get control back
        self.assertEqual(sorted(unwound), [1, 2])

    def test_return_exceptions_collects_results(self):

        async def boom():
            raise ValueError('boom')

        async def ok(n):
            return n

        results = self._await(
            run_limited([ok(1), boom(), ok(3)], 4, return_exceptions=True))
        self.assertEqual([results[0], results[2]], [1, 3])
        self.assertIsInstance(results[1], ValueError)

    def test_concurrent_failure_is_logged_not_dropped(self):
        # Both fail before we regain control, but only one can propagate; the
        # other would otherwise vanish -- gather marks it retrieved for us.
        async def boom(exc):
            raise exc

        with self.assertLogs(level='WARNING') as logs:
            self.assertRaises((ValueError, KeyError), self._await,
                              run_limited([boom(ValueError('a')),
                                           boom(KeyError('b'))], 4))
        discarded = [ln for ln in logs.output if 'run_limited: discarding' in ln]
        self.assertEqual(len(discarded), 1, logs.output)

    def test_caller_cancellation_waits_for_cleanup(self):
        unwound = []

        async def slow(n, unwind_delay):
            try:
                await asyncio.sleep(30)
            except asyncio.CancelledError:
                await asyncio.sleep(unwind_delay)  # e.g. closing a client
                unwound.append(n)
                raise

        async def body():
            coros = [slow(0, 0.001), slow(1, 0.02), slow(2, 0.04)]
            task = asyncio.ensure_future(run_limited(coros, 8))
            await asyncio.sleep(0.01)  # let them all start
            task.cancel()
            with self.assertRaises(asyncio.CancelledError):
                await task

        self._await(body())
        # every child unwound, including the slow ones: we neither returned
        # early nor cut their cleanup short with a second cancel
        self.assertEqual(sorted(unwound), [0, 1, 2])

    def test_closes_coroutines_still_queued(self):
        ran = []

        async def hog():
            await asyncio.sleep(30)

        async def never():
            ran.append(True)

        # limit of 1, so the second coroutine is still parked on the semaphore
        queued = never()

        async def body():
            task = asyncio.ensure_future(run_limited([hog(), queued], 1))
            await asyncio.sleep(0.01)
            task.cancel()
            with self.assertRaises(asyncio.CancelledError):
                await task

        self._await(body())
        self.assertEqual(ran, [])
        # closed, so it will not warn about never being awaited
        self.assertIsNone(queued.cr_frame)

    def test_child_cancellation_does_not_wait_on_siblings(self):
        # A child raising CancelledError on its own is not the caller being
        # cancelled: gather has not touched the siblings, so we must cancel
        # them rather than block until they finish by themselves.
        unwound = []

        async def self_cancelling():
            raise asyncio.CancelledError

        async def slow(n):
            try:
                await asyncio.sleep(30)
            except asyncio.CancelledError:
                unwound.append(n)
                raise

        async def body():
            task = asyncio.ensure_future(
                run_limited([self_cancelling(), slow(1)], 4))
            done, pending = await asyncio.wait([task], timeout=5)
            self.assertEqual(pending, set(), 'hung waiting for a sibling')
            self.assertRaises(asyncio.CancelledError, task.result)

        self._await(body())
        self.assertEqual(unwound, [1])


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