Unit tests rewrite