Improve tests, tidy up a little.