Fix logic in model test