Fix a couple of tests

Fix a couple of tests that were broken by commit 7e068a3.