Heroku fix
Updated tests