Unit tests fixes