Added test and removed bugs (tests are awesome!)