Adding tests and some changes