Add tests, fix minor errors