Update test suite, refactor