Code review fixes in tests