Fixed a stupid bug & added tests