[BUG] Fix test flake from 5283. (#5287)

## Description of changes

Previous PR introduced a flake.  Fix it.

## Test plan

Local for failing test + CI

## Migration plan

N/A

## Observability plan

N/A

## Documentation Changes

N/A