Adding tests, and fixing code according to tests