Unit tests cleanup