import asyncio


def wrap(fn, fn_name, called):
    """Wrap fn to add fn_name to called on invocation.
    """
    async def async_wrapper(*args, **kwargs):
        await fn(*args, **kwargs)
        called.add(fn_name)
    def wrapper(*args, **kwargs):
        fn(*args, **kwargs)
        called.add(fn_name)
    return async_wrapper if asyncio.iscoroutinefunction(fn) else wrapper