Minor code cleanup in unit tests