Fix some tests broken by previous commit