Added more tests and minor fixes