Small fix and update of tests