Fixes and more tests (with less code!)