[Tests] Fixing broken unit tests