Merge "Fix unit test when the flag is enabled" into main