Adding tests and polish