Fixing and adding to test suite