Minor code refactor and cleanup for tests.