Fixing broken unit tests