Test code refactor