Cleanup and some tests