Adding tests and fixes