Merge "Add flag guard and fix test" into main