One of the patterns I like is using a context manager to make sure bringup and cleanup are called for each test. Let’s have a look at context managers.

Context Manager and Async Context Manager Link to heading

I will just reference the official PEP 343

This PEP adds a new statement “with” to the Python language to make it possible to factor out standard uses of try/finally statements.

In this PEP, context managers provide enter() and exit() methods that are invoked on entry to and exit from the body of the with statement.

Essentially, using with we can ensure __enter__ and __exit__ are called to do the initialization and cleanup code.

with open(filename) as f:
    ...read data from f...

Then in PEP 492, they introduced the awaitable version of a context manager:

An asynchronous context manager is a context manager that is able to suspend execution in its enter and exit methods.

To make this possible, a new protocol for asynchronous context managers is proposed. Two new magic methods are added: aenter and aexit. Both must return an awaitable.

class AsyncContextManager:
    async def __aenter__(self):
        await log('entering context')

    async def __aexit__(self, exc_type, exc, tb):
        await log('exiting context')

cocotb Context Manager Link to heading

It’s always good to have self-contained reset, bringup, and teardown methods so we can just call them in __aenter__ and __aexit__.

That said, there is a simpler way of doing this (using __init__ and finally), but the context manager is too elegant to give up.

import cocotb
from cocotb.clock import Clock
from cocotb.triggers import ClockCycles

class TestBench:
    def __init__(self, dut):
        self.dut = dut
        self.clock_task = None

    async def reset(self):
        cocotb.log.info("Applying reset")
        self.dut.rst_n.value = 0
        await ClockCycles(self.dut.clk, 5)
        
        cocotb.log.info("Releasing reset")
        self.dut.rst_n.value = 1
        await ClockCycles(self.dut.clk, 2)

    async def bringup(self):
        cocotb.log.info("Bringup starting")
        
        # Start the clock
        self.clock_task = cocotb.start_soon(
            Clock(self.dut.clk, 10, unit='ns').start()
        )

        # Reset the DUT
        await self.reset()
        cocotb.log.info("Bringup completed")

    async def teardown(self):
        cocotb.log.info("Teardown starting")

        # Do stuff to leave the DUT in a known state
        self.dut.data_in.value = 0
        await self.reset()

        # Stop the clock
        self.clock_task.cancel()

        cocotb.log.info("Teardown completed")

    async def __aenter__(self):
        await self.bringup()
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.teardown()
        return False
     
@cocotb.test()
async def test_async_context_manager(dut):
    cocotb.log.info("Test started")

    async with TestBench(dut) as tb:
        await ClockCycles(dut.clk, 10)
        dut.data_in.value = 0xDEADBEEF
        cocotb.log.info(f"Set data_in = 0x{int(dut.data_in.value):X}")
        await ClockCycles(dut.clk, 5)
        
    cocotb.log.info("Test Completed")

     0.00ns INFO     test                               Test started
     0.00ns INFO     test                               Bringup starting
     0.00ns INFO     test                               Applying reset
    40.00ns INFO     test                               Releasing reset
    60.00ns INFO     test                               Bringup completed
   160.00ns INFO     test                               Set data_in = 0xDEADBEEF
   210.00ns INFO     test                               Teardown starting
   210.00ns INFO     test                               Applying reset
   260.00ns INFO     test                               Releasing reset
   280.00ns INFO     test                               Teardown completed
   280.00ns INFO     test                               Test Completed