Fixing tests!