Fixed failing unit test.

Fixed failing unit test following previous commit.