Fix unit tests
adressed @romainneutron's comments