More unit tests fixes