Add tests and fix a critical bug