Adding tests and fixing issues