Deterministic Simulation Testing with Python async/await

sleep-sim.py

import asyncio
from heapq import heappop, heappush

async def sweet_sleep(sleep):
    await sleep(1)
    return 'hey'

def main():
    clock_ms = 0
    time_ms = lambda: clock_ms

    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

    sleep_queue = []
    async def sleep(seconds):
        future = loop.create_future()
        heappush(sleep_queue, (time_ms() + seconds * 1000, future))
        await future

    task = loop.create_task(sweet_sleep(sleep))
    def verify_result(task):
        assert time_ms() == 1000
        assert task.result() == 'hey'
    task.add_done_callback(verify_result)

    while sleep_queue or asyncio.all_tasks(loop):
        if sleep_queue:
            next_time, next_future = min(sleep_queue)
            if next_time <= time_ms():
                heappop(sleep_queue)
                next_future.set_result(None)
        loop.run_until_complete(asyncio.sleep(0))
        clock_ms += 1

    assert time_ms() == 1001
    loop.stop()
    loop.close()

if __name__ == '__main__':
    import sys

    sys.exit(main())

sleep.py

import asyncio

async def sweet_sleep(sleep):
    print('1:starting sweet sleep')
    await sleep(1)
    print('5:woke up from sweet sleep')
    return 'happy'

# create a new event loop for the current thread
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)

sleep_future = loop.create_future()
requested_sleep = 0
async def sleep(duration):
    global requested_sleep

    print(f'2:requested to sleep for {duration}s')
    requested_sleep = duration
    await sleep_future

task = loop.create_task(sweet_sleep(sleep))
print('0:task created')

# allow other coroutines to be run on the current thread
loop.run_until_complete(asyncio.sleep(0))

print('3:back to main thread')
assert requested_sleep == 1
sleep_future.set_result(None)
print('4:unblocked the future')

# run other coroutines
loop.run_until_complete(asyncio.sleep(0))
print(f'6:task result: {task.result()}')