Fixed failing test

Fixed unit test