Fix tests.

Hey! This looks really awesome.

I fixed a couple of tests that were broken, and cleanup up some duplication.