Added more tests and fixed docs