Fixed failing tests from model refactor