Code cleanup for existing generator unit tests