Try to fix up tests