Refactor unit tests to use our own functionality to determine whether or not to retry