More tests, fix bugs