# AUTOGENERATED - modify shared_anyscale_util in root directory to make changes
import asyncio

import pytest

from shared_anyscale_utils.utils.asyncio import chain_future, create_future


@pytest.mark.parametrize("result", [42, Exception("failed!")])
async def test_chain_future_completes(result):
    one = create_future()
    other = create_future()

    chain_future(one, other)

    if isinstance(result, Exception):
        one.set_exception(result)
    else:
        one.set_result(result)

    await asyncio.gather(other, return_exceptions=True)

    if isinstance(result, Exception):
        assert str(other.exception()) == str(result)
    else:
        assert other.result() == 42


async def test_chain_future_cancellation():
    one = create_future()
    other = create_future()

    chain_future(one, other)

    one.cancel()

    # NOTE: We're using `gather` to retrieve `CancelledError` since `future.exception` will
    #       raise it instead of returning it
    e = (await asyncio.gather(other, return_exceptions=True))[0]

    assert type(e) == asyncio.CancelledError
