Small refactor of cleanup for test