Refactor flaky test

[gabegma]

Description:

• This test sometimes fails, and I don't understand why. I cleaned it up, but I'm not sure if my code is equivalent and could improve anything. Opening this PR to brainstorm. Any ides @JosephMarinier @Dref360 ?
  

Checklist:

You should check all boxes before the PR is ready. If a box does not apply, check it to acknowledge it.

• ISSUE NUMBER. You linked the issue number (Ex: Resolve #XXX).
• PRE-COMMIT. You ran pre-commit on all commits, or else, you
  ran pre-commit run --all-files at the end.
• USER CHANGES. The changes are added to CHANGELOG.md and the documentation, if they impact
  our users.
• DEV CHANGES.
  • Update the documentation if this PR changes how to develop/launch on the app.
  • Update the README files and our wiki for any big design decisions, if relevant.
  • Add unit tests, docstrings, typing and comments for complex sections.

[JosephMarinier]

As you said, I'm not sure we solved the problem, but we made the code easier to understand and with the added assertion print we should be able to understand the error next time it happens. 👍