Fixing up first response, added tests!