Rewrite tests to use afterEach