Fix unit tests.

Make unit tests more robust (and fix a couple of logic bugs) to make the
tests deterministic (and pass all the time).