Adding tests and fixing errors