Fix issues in tests