Small test refactor for performance

Small performance improvement.

Closes #724