Cleanup + fixed the tests because I accidentally the tests..