Revert "Revert "unit test fix""

This reverts commit d1a9bcb0a17f28c4487b67e8d12ee056757e338a.