Fix unit tests from previous commit

Change-Id: I5e4459d83159db72dd70e78aca0fe75658f4e77a