Unit tests improvements