Fixed it a little bit, pimped the tests also