Removes after test setup and adds small refactor.