More tests and cleanup