Minor test code cleanup