Fix unit tests following code layout update