Fixing up tests