diff --git a/tests/test_rapids-conda-retry.py b/tests/test_rapids-conda-retry.py new file mode 100644 index 0000000..40fb600 --- /dev/null +++ b/tests/test_rapids-conda-retry.py @@ -0,0 +1,85 @@ +import os +import stat +import subprocess +import tempfile +from os import path + +import pytest + +TOOLS_DIR = path.join(path.dirname(path.realpath(__file__)), "..", "tools") + +FAKE_CONDA = """#!/bin/bash +count=$(cat "$COUNT_FILE" 2>/dev/null || echo 0) +echo $((count + 1)) > "$COUNT_FILE" +if (( count < FAIL_TIMES )); then + echo "$FAIL_OUTPUT" + exit 1 +fi +exit 0 +""" + + +def run_rapids_conda_retry(fail_output, fail_times): + with tempfile.TemporaryDirectory() as d: + fake_conda = path.join(d, "fake-conda") + count_file = path.join(d, "count") + with open(fake_conda, "w") as f: + f.write(FAKE_CONDA) + os.chmod(fake_conda, os.stat(fake_conda).st_mode | stat.S_IXUSR) + + env = { + **os.environ, + "PATH": f"{TOOLS_DIR}:{os.environ['PATH']}", + "RAPIDS_CONDA_EXE": fake_conda, + "RAPIDS_CONDA_RETRY_MAX": "2", + "RAPIDS_CONDA_RETRY_SLEEP": "0", + "COUNT_FILE": count_file, + "FAIL_TIMES": str(fail_times), + "FAIL_OUTPUT": fail_output, + } + process = subprocess.run( + [path.join(TOOLS_DIR, "rapids-conda-retry"), "install", "pkg"], + env=env, + text=True, + capture_output=True, + ) + with open(count_file) as f: + return process, int(f.read()) + + +@pytest.mark.parametrize( + "fail_output", + [ + "requests.exceptions.HTTPError: 500 Server Error: Internal Server Error for url: https://conda.anaconda.org/x", + "requests.exceptions.HTTPError: 502 Server Error: Bad Gateway for url: https://conda.anaconda.org/x", + "requests.exceptions.HTTPError: 503 Server Error: Service Unavailable for url: https://conda.anaconda.org/x", + ], +) +def test_retries_on_5xx_http_error(fail_output): + process, calls = run_rapids_conda_retry(fail_output, fail_times=1) + assert process.returncode == 0 + assert calls == 2 + assert "Retrying, found 'HTTPError:' in output..." in process.stderr + + +def test_gives_up_after_max_retries_on_http_error(): + output = "requests.exceptions.HTTPError: 502 Server Error: Bad Gateway for url: https://conda.anaconda.org/x" + process, calls = run_rapids_conda_retry(output, fail_times=100) + assert process.returncode == 1 + assert calls == 3 + assert "reached max retries" in process.stderr + + +@pytest.mark.parametrize("status", ["404 Client Error: Not Found", "429 Client Error: Too Many Requests"]) +def test_retries_on_4xx_http_error(status): + output = f"requests.exceptions.HTTPError: {status} for url: https://conda.anaconda.org/x" + process, calls = run_rapids_conda_retry(output, fail_times=1) + assert process.returncode == 0 + assert calls == 2 + + +def test_does_not_retry_on_unrelated_error(): + process, calls = run_rapids_conda_retry("PackagesNotFoundError: pkg", fail_times=1) + assert process.returncode == 1 + assert calls == 1 + assert "no retryable" in process.stderr diff --git a/tools/rapids-conda-retry b/tools/rapids-conda-retry index 6b70d1a..57a0f61 100755 --- a/tools/rapids-conda-retry +++ b/tools/rapids-conda-retry @@ -95,6 +95,9 @@ function runConda { elif grep -q CondaHTTPError: "${outfile}"; then retryingMsg="Retrying, found 'CondaHTTPError:' in output..." needToRetry=1 + elif grep -q HTTPError: "${outfile}"; then + retryingMsg="Retrying, found 'HTTPError:' in output..." + needToRetry=1 elif grep -q CondaMultiError: "${outfile}"; then retryingMsg="Retrying after cleaning tarball cache, found 'CondaMultiError:' in output..." needToRetry=1 @@ -167,6 +170,7 @@ function runConda { 'Failed writing received data to disk', \ 'File not valid: file size doesn't match expectation', \ 'File not valid: SHA256 sum doesn't match expectation', \ +'HTTPError:', \ 'JSONDecodeError:', \ 'Multi-download failed', \ 'Response ended prematurely', \