Fix class tests