Fixing stuff with tests