Adding tests and fixing bugs