Fixing existing tests