Fix up some of the unit testing code.