Adding tests and fixing a minor bug