Fix tests

Thanks @umaar for noticing :)