Fix unit tests broken by model registry changes