diff --git a/core/network_utils.py b/core/network_utils.py index ffd8d71..2d62f9f 100644 --- a/core/network_utils.py +++ b/core/network_utils.py @@ -208,8 +208,9 @@ def request_source_response( timeout: int = 20, max_attempts: int = 2, log_name: str = "Source API", + raise_on_error: bool = False, ) -> Optional[requests.Response]: - """Perform a GET/POST request with retry logic for rate limit status (HTTP 429).""" + """Perform a source request with optional transport-error propagation.""" import time from .log_system import create_module_logger @@ -226,6 +227,8 @@ def request_source_response( response = requests.get(url, params=request_params, headers=headers, timeout=timeout) except Exception as e: log.warning(f"{log_name} request failed: url={url}, error={e}") + if raise_on_error: + raise return None if response.status_code != 429 or attempt == max_attempts - 1: diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index b7d1f9f..e882648 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -11,6 +11,8 @@ import time from typing import Any, Callable, Dict, List, Optional, Tuple from urllib.parse import urljoin, urlparse +import requests + from ..log_system import create_module_logger from ..matcher import ( base_model_matches as _base_model_matches, @@ -96,6 +98,37 @@ REQUEST_HEADERS = { class CivArchiveSearchError(Exception): """Raised when CivArchive cannot complete a search request.""" + def __init__( + self, + message: str, + *, + code: str = "provider_unavailable", + http_status: Optional[int] = None, + retryable: bool = True, + ) -> None: + super().__init__(message) + self.code = code + self.http_status = http_status + self.retryable = retryable + + +def _classify_civarchive_http_status(status_code: int) -> Dict[str, Any]: + """Classify an HTTP response for user-facing provider status handling.""" + try: + normalized_status = int(status_code) + except (TypeError, ValueError): + return {"code": "provider_unavailable", "retryable": True} + + if normalized_status == 408: + return {"code": "timeout", "retryable": True} + if normalized_status == 429: + return {"code": "rate_limited", "retryable": True} + if normalized_status == 404: + return {"code": "not_found", "retryable": False} + if 400 <= normalized_status < 500: + return {"code": "provider_rejected", "retryable": False} + return {"code": "provider_unavailable", "retryable": True} + def clear_search_cache(): """Clear cached CivArchive search results.""" @@ -433,22 +466,47 @@ def _search_page( # (especially VAEs) are frequently misclassified (e.g. as Checkpoint or null type) in their database. # We rely on name confidence sorting instead. - response = request_source_response( - f"{CIVARCHIVE_BASE_URL}/search", - method="GET", - headers=REQUEST_HEADERS, - params=params, - timeout=timeout, - log_name="CivArchive search", - ) + try: + response = request_source_response( + f"{CIVARCHIVE_BASE_URL}/search", + method="GET", + headers=REQUEST_HEADERS, + params=params, + timeout=timeout, + log_name="CivArchive search", + raise_on_error=True, + ) + except requests.Timeout as error: + raise CivArchiveSearchError( + "CivArchive search request timed out", + code="timeout", + retryable=True, + ) from error + except requests.RequestException as error: + raise CivArchiveSearchError( + "CivArchive search request failed (network error)", + code="network_error", + retryable=True, + ) from error + if response is None: - raise CivArchiveSearchError("CivArchive search request failed (network error or timeout)") + raise CivArchiveSearchError( + "CivArchive search request failed (network error or timeout)", + code="network_error", + retryable=True, + ) if response.status_code != 200: + failure = _classify_civarchive_http_status(response.status_code) log.warning( f"CivArchive search returned {response.status_code}: query={query}" ) - raise CivArchiveSearchError(f"HTTP {response.status_code}") + raise CivArchiveSearchError( + f"HTTP {response.status_code}", + code=failure["code"], + http_status=response.status_code, + retryable=failure["retryable"], + ) next_data = _extract_next_data(response.text) results = ( diff --git a/tests/test_civarchive_failure_contract.py b/tests/test_civarchive_failure_contract.py index d01813b..7de8ea7 100644 --- a/tests/test_civarchive_failure_contract.py +++ b/tests/test_civarchive_failure_contract.py @@ -2,24 +2,84 @@ from types import SimpleNamespace from unittest.mock import patch import pytest +import requests +from core import network_utils from core.sources import civarchive +@pytest.mark.parametrize( + ("status_code", "expected_code", "retryable"), + [ + (408, "timeout", True), + (429, "rate_limited", True), + (500, "provider_unavailable", True), + (502, "provider_unavailable", True), + (522, "provider_unavailable", True), + (404, "not_found", False), + (403, "provider_rejected", False), + ], +) +def test_civarchive_http_failures_are_classified(status_code, expected_code, retryable): + assert civarchive._classify_civarchive_http_status(status_code) == { + "code": expected_code, + "retryable": retryable, + } + + @pytest.mark.parametrize("status_code", [502, 522]) def test_civarchive_search_preserves_http_failure_message(status_code): response = SimpleNamespace(status_code=status_code, text="") - with ( - patch.object(civarchive, "request_source_response", return_value=response), - pytest.raises(civarchive.CivArchiveSearchError, match=f"HTTP {status_code}"), - ): + with patch.object(civarchive, "request_source_response", return_value=response), pytest.raises( + civarchive.CivArchiveSearchError, match=f"HTTP {status_code}" + ) as raised: civarchive._search_page("model.safetensors") + assert raised.value.code == "provider_unavailable" + assert raised.value.http_status == status_code + assert raised.value.retryable is True + def test_civarchive_search_preserves_network_failure_message(): - with ( - patch.object(civarchive, "request_source_response", return_value=None), - pytest.raises(civarchive.CivArchiveSearchError, match="network error or timeout"), - ): + with patch.object(civarchive, "request_source_response", return_value=None), pytest.raises( + civarchive.CivArchiveSearchError, match="network error or timeout" + ) as raised: civarchive._search_page("model.safetensors") + + assert raised.value.code == "network_error" + assert raised.value.http_status is None + assert raised.value.retryable is True + + +@pytest.mark.parametrize( + ("transport_error", "expected_code", "expected_message"), + [ + (requests.Timeout("read timeout"), "timeout", "timed out"), + (requests.ConnectionError("connection reset"), "network_error", "network error"), + ], +) +def test_civarchive_search_classifies_transport_failures( + transport_error, + expected_code, + expected_message, +): + with patch.object( + civarchive, + "request_source_response", + side_effect=transport_error, + ), pytest.raises(civarchive.CivArchiveSearchError, match=expected_message) as raised: + civarchive._search_page("model.safetensors") + + assert raised.value.code == expected_code + assert raised.value.retryable is True + + +def test_request_source_response_can_propagate_transport_errors(): + with patch.object(network_utils.requests, "get", side_effect=requests.Timeout("timeout")), pytest.raises( + requests.Timeout, match="timeout" + ): + network_utils.request_source_response( + "https://example.test/search", + raise_on_error=True, + )