Fix stuff
Add test
Fix code a bit