Refactor model test code