Fix unit tests

Change-Id: I4329bb8b101c4797737fa879d7eeb9be57effdc7