Fixing up test suite