Fix models
Update tests