Rewriting test code to simplify it.