added a few more tests and cleaned up a little bit