Added some unit tests, fixed some bugs!