Fix up some tests