Fix more tests.

Will be squashed upon merge