Python contextlib.asynccontextmanager() Examples

def test_basic():
    exit_count = 0

    @contextlib.asynccontextmanager
    async def ctx(v):
        nonlocal exit_count
        await asyncio.sleep(0)
        yield v
        await asyncio.sleep(0)
        exit_count += 1

    group = AsyncContextGroup([ctx(i) for i in range(3)])
    async with group as yielded_values:
        assert yielded_values == tuple(range(3))

    assert exit_count == 3 

def test_exception_entering_context():
    exit_count = 0

    @contextlib.asynccontextmanager
    async def ctx(should_raise=False):
        nonlocal exit_count
        await asyncio.sleep(0)
        if should_raise:
            raise ValueError()
        try:
            yield
        finally:
            await asyncio.sleep(0)
            exit_count += 1

    group = AsyncContextGroup([ctx(), ctx(True), ctx()])
    with pytest.raises(ValueError):
        async with group:
            # the body of the with block should never execute if an exception is raised when
            # entering the context group.
            assert False  # noqa: B011

    # One of our contexts was not entered so we didn't exit it either.
    assert exit_count == 2 

def test_exception_exiting():
    exit_count = 0

    @contextlib.asynccontextmanager
    async def ctx(should_raise=False):
        nonlocal exit_count
        await asyncio.sleep(0)
        try:
            yield
        finally:
            exit_count += 1
            if should_raise:
                raise ValueError()

    group = AsyncContextGroup([ctx(), ctx(True), ctx(True)])
    with pytest.raises(MultiError) as exc_info:
        async with group:
            pass

    exc = exc_info.value
    assert len(exc.exceptions) == 2
    assert isinstance(exc.exceptions[0], ValueError)
    assert isinstance(exc.exceptions[1], ValueError)
    assert exit_count == 3 

def test_exception_inside_context_block():
    exit_count = 0

    async def f(should_raise):
        await asyncio.sleep(0)
        if should_raise:
            raise ValueError()

    @contextlib.asynccontextmanager
    async def ctx(should_raise=False):
        nonlocal exit_count
        await asyncio.sleep(0)
        try:
            yield f(should_raise)
        finally:
            await asyncio.sleep(0)
            exit_count += 1

    group = AsyncContextGroup([ctx(), ctx(True), ctx()])
    with pytest.raises(ValueError):
        async with group as awaitables:
            empty1, exception, empty2 = await asyncio.gather(*awaitables, return_exceptions=True)
            assert empty1 is None
            assert empty2 is None
            raise exception

    assert exit_count == 3 

def test_enter_context(self):
        state = None

        @contextlib.contextmanager
        def ctxtest():
            nonlocal state
            state = "before"
            yield state
            state = "after"

        @contextlib.asynccontextmanager
        async def actxtest():
            nonlocal state
            state = "before2"
            yield state
            state = "after2"

        async with await s_base.Base.anit() as base:
            await base.enter_context(ctxtest())
            self.eq("before", state)

        self.eq("after", state)

        async with await s_base.Base.anit() as base:
            await base.enter_context(actxtest())
            self.eq("before2", state)

        self.eq("after2", state) 

def fake_beam_syncer(chain, event_bus):
    @contextlib.asynccontextmanager
    async def fake_beam_sync(removed_nodes: Dict):
        # beam sync starts, it fetches requested nodes from remote peers

        def replace_missing_node(missing_node_hash):
            if missing_node_hash not in removed_nodes:
                raise Exception(f'An unexpected node was requested: {missing_node_hash}')
            chain.chaindb.db[missing_node_hash] = removed_nodes.pop(missing_node_hash)

        async def collect_accounts(event: CollectMissingAccount):
            replace_missing_node(event.missing_node_hash)
            await event_bus.broadcast(
                MissingAccountResult(1), event.broadcast_config()
            )
        accounts_sub = event_bus.subscribe(CollectMissingAccount, collect_accounts)

        async def collect_bytecodes(event: CollectMissingBytecode):
            replace_missing_node(event.bytecode_hash)
            await event_bus.broadcast(
                MissingBytecodeResult(), event.broadcast_config()
            )
        bytecode_sub = event_bus.subscribe(CollectMissingBytecode, collect_bytecodes)

        async def collect_storage(event: CollectMissingStorage):
            replace_missing_node(event.missing_node_hash)
            await event_bus.broadcast(
                MissingStorageResult(1), event.broadcast_config()
            )
        storage_sub = event_bus.subscribe(CollectMissingStorage, collect_storage)

        await event_bus.wait_until_any_endpoint_subscribed_to(CollectMissingAccount)
        await event_bus.wait_until_any_endpoint_subscribed_to(CollectMissingBytecode)
        await event_bus.wait_until_any_endpoint_subscribed_to(CollectMissingStorage)

        try:
            yield
        finally:
            accounts_sub.unsubscribe()
            bytecode_sub.unsubscribe()
            storage_sub.unsubscribe()

    return fake_beam_sync


# Test that eth_getBalance works during beam sync