Fixed some issues and added a test (untested)