Alter a ugly code on unit.rb and refactor this