Fix migration tests