First attempt to fix tests