Little changes into unit testing