Improve a couple of tests