import cocotb
from cocotb.clock import Clock
from cocotb.triggers import FallingEdge, ReadOnly, RisingEdge


@cocotb.test()
async def counter_counts_when_enabled(dut):
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())

    dut.rst.value = 1
    dut.en.value = 0
    await RisingEdge(dut.clk)
    await RisingEdge(dut.clk)

    await FallingEdge(dut.clk)
    dut.rst.value = 0
    dut.en.value = 1

    for expected in range(1, 6):
        await RisingEdge(dut.clk)
        await ReadOnly()
        assert dut.count.value.to_unsigned() == expected

    await FallingEdge(dut.clk)
    dut.en.value = 0
    await ReadOnly()
    held = dut.count.value.to_unsigned()
    await RisingEdge(dut.clk)
    await ReadOnly()
    assert dut.count.value.to_unsigned() == held
    await RisingEdge(dut.clk)
    await ReadOnly()
    assert dut.count.value.to_unsigned() == held
