Some cleanup
Added some more tests