From d295f238cdfd195bc46cf79acfd382522d3f0e21 Mon Sep 17 00:00:00 2001 From: Greg Lucas Date: Mon, 9 Feb 2026 09:30:10 -0700 Subject: [PATCH] MNT: Add retry logic and better exception handling --- imap_data_access/io.py | 8 ++++++-- tests/test_io.py | 9 +++------ 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/imap_data_access/io.py b/imap_data_access/io.py index fc6c3eb..e35e499 100644 --- a/imap_data_access/io.py +++ b/imap_data_access/io.py @@ -24,6 +24,9 @@ class IMAPDataAccessError(Exception): pass +_RETRY_ADAPTER = requests.adapters.HTTPAdapter(max_retries=3) + + @contextlib.contextmanager def _make_request(request: requests.PreparedRequest): """Get the response from a URL request using the requests library. @@ -44,6 +47,7 @@ def _make_request(request: requests.PreparedRequest): ) try: with requests.Session() as session: + session.mount("https://", _RETRY_ADAPTER) response = session.send(request) response.raise_for_status() yield response @@ -52,8 +56,8 @@ def _make_request(request: requests.PreparedRequest): error_msg = f"{e.response.status_code} {e.response.reason}: {e.response.text}" raise IMAPDataAccessError(error_msg) from e except requests.exceptions.RequestException as e: - error_msg = f"{e.response.status_code} {e.response.reason}: {e.response.text}" - raise IMAPDataAccessError(error_msg) from e + # Handle cases where response may not exist (connection errors, timeouts, etc.) + raise IMAPDataAccessError(str(e)) from e def _get_base_url() -> str: diff --git a/tests/test_io.py b/tests/test_io.py index 952eacd..926cd6d 100644 --- a/tests/test_io.py +++ b/tests/test_io.py @@ -97,15 +97,12 @@ def test_request_errors(mock_send_request): with pytest.raises(imap_data_access.io.IMAPDataAccessError, match="404 Not Found"): imap_data_access.download(test_science_path) - # Set up the mock to raise a RequestException - mock_response.status_code = 400 - mock_response.reason = "Request failed" - mock_response.text = "" + # Set up the mock to raise a RequestException with a response mock_send_request.side_effect = requests.exceptions.RequestException( - response=mock_response + "connection error" ) with pytest.raises( - imap_data_access.io.IMAPDataAccessError, match="400 Request failed" + imap_data_access.io.IMAPDataAccessError, match="connection error" ): imap_data_access.download(test_science_path)