Fix bug in the testing code.

See https://piazza.com/class/hzb83ip5yb824g?cid=120