Put back code so tests keep working