Fix unit tests.

Change-Id: If0edb209ff80cff49890e4468de15ea987664de6