Refactor test map code to be more modifiable