Fixed up the test cases, and made look a bit sexier