Back-pressure in chiseltest

ExtendedValidDriver.scala

// I cut and paste this from a bigger project, so there could be missing imports/etc
import chisel3._
import chisel3.util._
import chisel3.experimental.{DataMirror, Direction}
import chiseltest._
import chiseltest.testableData
import scala.util.Random
import scala.language.implicitConversions

trait Stalls {
  def get: Int = 0
}

case object NoStalls extends Stalls

case class RandomStalls(minStalls: Int, maxStalls: Int) extends Stalls {
  require(minStalls <= maxStalls, s"maxStalls ($maxStalls) needs to be greater than or equal to minStalls ($minStalls)")

  override def get: Int = minStalls + Random.nextInt(maxStalls - minStalls + 1)
}
case class FixedStalls(stalls: Int) extends Stalls {
  override def get: Int = stalls
}

class ExtendedDecoupledDriver[T <: Data](x: ReadyValidIO[T], stalls: Stalls = NoStalls) extends DecoupledDriver(x) {

  var cycles = 0

  override def enqueue(data: T): Unit = timescope {
    sourceValid()
    x.bits.poke(data)
    fork
      .withRegion(Monitor) {
        while (x.ready.peek().litToBoolean == false) {
          getSourceClock.step(1)
        }
      }
      .joinAndStep(getSourceClock)
  }

  override def expectDequeue(data: T): Unit = timescope {
    sinkReady()
    fork
      .withRegion(Monitor) {
        waitForValid()
        // x.valid.expect(true.B)
        x.bits.expect(data)
      }
      .joinAndStep(getSinkClock)
  }

  def withStalls(_stalls: Stalls): ExtendedDecoupledDriver[T] = {
    new ExtendedDecoupledDriver(x, _stalls)
  }

  override def waitForValid(): Unit = {
    val clock = getSinkClock
    while (!x.valid.peek().litToBoolean) {
      clock.step(1)
      cycles += 1
    }
  }

  def waitForReadyAndStep(): Unit = {
    fork
      .withRegion(Monitor) {
        while (!x.ready.peek().litToBoolean) {
          getSourceClock.step(1)
          cycles += 1
        }
      }
      .joinAndStep(getSourceClock)
    cycles += 1
  }

  def sinkReady(): Unit = {
    for (_ <- 0 until stalls.get) {
      x.ready.poke(false.B)
      getSinkClock.step(1)
      cycles += 1
    }
    x.ready.poke(true.B)
  }

  def sourceValid(): Unit = {
    for (_ <- 0 until stalls.get) {
      x.valid.poke(false.B)
      getSourceClock.step(1)
      cycles += 1
    }
    x.valid.poke(true.B)
  }

  def requiresSink() = {
    require(
      DataMirror.directionOf(x.valid) == Direction.Output &&
        DataMirror.directionOf(x.bits) == Direction.Output &&
        DataMirror.directionOf(x.ready) == Direction.Input,
      "a sink (output) port is required"
    )
  }
}

object ExtendedDecoupledDriver {
  implicit def decoupledToExtendedDecoupledDriver[T <: Data](x: ReadyValidIO[T]): ExtendedDecoupledDriver[T] =
    new ExtendedDecoupledDriver(x)
}