Improved UI a bit and updated test coverage accordingly