Fixed and added a unittest, thanks