Refactor retry function and add unit tests.