Unit tests refactor