Fixing tests

fix tests