# Copyright 2018 Gentoo Foundation
# Distributed under the terms of the GNU General Public License v2

from portage.util.futures import asyncio
from portage.util.futures.compat_coroutine import (
	coroutine,
	coroutine_return,
)
from portage.util.futures._sync_decorator import _sync_decorator, _sync_methods
from portage.tests import TestCase


class CompatCoroutineTestCase(TestCase):

	def test_returning_coroutine(self):
		@coroutine
		def returning_coroutine():
			yield asyncio.sleep(0)
			coroutine_return('success')

		self.assertEqual('success',
			asyncio.get_event_loop().run_until_complete(returning_coroutine()))

	def test_raising_coroutine(self):

		class TestException(Exception):
			pass

		@coroutine
		def raising_coroutine():
			yield asyncio.sleep(0)
			raise TestException('exception')

		self.assertRaises(TestException,
			asyncio.get_event_loop().run_until_complete, raising_coroutine())

	def test_catching_coroutine(self):

		class TestException(Exception):
			pass

		@coroutine
		def catching_coroutine(loop=None):
			loop = asyncio._wrap_loop(loop)
			future = loop.create_future()
			loop.call_soon(future.set_exception, TestException('exception'))
			try:
				yield future
			except TestException:
				self.assertTrue(True)
			else:
				self.assertTrue(False)
			coroutine_return('success')

		loop = asyncio.get_event_loop()
		self.assertEqual('success',
			loop.run_until_complete(catching_coroutine(loop=loop)))

	def test_cancelled_coroutine(self):
		"""
		Verify that a coroutine can handle (and reraise) asyncio.CancelledError
		in order to perform any necessary cleanup. Note that the
		asyncio.CancelledError will only be thrown in the coroutine if there's
		an opportunity (yield) before the generator raises StopIteration.
		"""
		loop = asyncio.get_event_loop()
		ready_for_exception = loop.create_future()
		exception_in_coroutine = loop.create_future()

		@coroutine
		def cancelled_coroutine(loop=None):
			loop = asyncio._wrap_loop(loop)
			while True:
				task = loop.create_future()
				try:
					ready_for_exception.set_result(None)
					yield task
				except BaseException as e:
					# Since python3.8, asyncio.CancelledError inherits
					# from BaseException.
					task.done() or task.cancel()
					exception_in_coroutine.set_exception(e)
					raise
				else:
					exception_in_coroutine.set_result(None)

		future = cancelled_coroutine(loop=loop)
		loop.run_until_complete(ready_for_exception)
		future.cancel()

		self.assertRaises(asyncio.CancelledError,
			loop.run_until_complete, future)

		self.assertRaises(asyncio.CancelledError,
			loop.run_until_complete, exception_in_coroutine)

	def test_cancelled_future(self):
		"""
		When a coroutine raises CancelledError, the coroutine's
		future is cancelled.
		"""

		@coroutine
		def cancelled_future_coroutine(loop=None):
			loop = asyncio._wrap_loop(loop)
			while True:
				future = loop.create_future()
				loop.call_soon(future.cancel)
				yield future

		loop = asyncio.get_event_loop()
		future = loop.run_until_complete(asyncio.wait([cancelled_future_coroutine()]))[0].pop()
		self.assertTrue(future.cancelled())

	def test_yield_expression_result(self):
		@coroutine
		def yield_expression_coroutine():
			for i in range(3):
				x = yield asyncio.sleep(0, result=i)
				self.assertEqual(x, i)

		asyncio.get_event_loop().run_until_complete(yield_expression_coroutine())

	def test_method_coroutine(self):

		class Cubby(object):

			_empty = object()

			def __init__(self, loop):
				self._loop = loop
				self._value = self._empty
				self._waiters = []

			def _notify(self):
				waiters = self._waiters
				self._waiters = []
				for waiter in waiters:
					waiter.cancelled() or waiter.set_result(None)

			def _wait(self):
				waiter = self._loop.create_future()
				self._waiters.append(waiter)
				return waiter

			@coroutine
			def read(self):
				while self._value is self._empty:
					yield self._wait()

				value = self._value
				self._value = self._empty
				self._notify()
				coroutine_return(value)

			@coroutine
			def write(self, value):
				while self._value is not self._empty:
					yield self._wait()

				self._value = value
				self._notify()

		@coroutine
		def writer_coroutine(cubby, values, sentinel):
			for value in values:
				yield cubby.write(value)
			yield cubby.write(sentinel)

		@coroutine
		def reader_coroutine(cubby, sentinel):
			results = []
			while True:
				result = yield cubby.read()
				if result == sentinel:
					break
				results.append(result)
			coroutine_return(results)

		loop = asyncio.get_event_loop()
		cubby = Cubby(loop)
		values = list(range(3))
		writer = asyncio.ensure_future(writer_coroutine(cubby, values, None), loop=loop)
		reader = asyncio.ensure_future(reader_coroutine(cubby, None), loop=loop)
		loop.run_until_complete(asyncio.wait([writer, reader]))

		self.assertEqual(reader.result(), values)

		# Test decoration of coroutine methods and functions for
		# synchronous usage, allowing coroutines to smoothly
		# blend with synchronous code.
		sync_cubby = _sync_methods(cubby, loop=loop)
		sync_reader = _sync_decorator(reader_coroutine, loop=loop)
		writer = asyncio.ensure_future(writer_coroutine(cubby, values, None), loop=loop)
		self.assertEqual(sync_reader(cubby, None), values)
		self.assertTrue(writer.done())

		for i in range(3):
			sync_cubby.write(i)
			self.assertEqual(sync_cubby.read(), i)
