Clean up test a bit