Fixing long-running flaky unit test

[mgrange1998]

Differential Revision: D91326985

[meta-codesync]

@mgrange1998 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D91326985.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 50.00000% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.71%. Comparing base (0c8cb23) to head (5539abe).

Files with missing lines	Patch %	Lines
ax/utils/testing/mock.py	50.00%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4808      +/-   ##
==========================================
- Coverage   96.71%   96.71%   -0.01%     
==========================================
  Files         587      587              
  Lines       61309    61315       +6     
==========================================
+ Hits        59297    59300       +3     
- Misses       2012     2015       +3     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.