diff --git a/mellea/formatters/granite/retrievers/util.py b/mellea/formatters/granite/retrievers/util.py index e5552fba7..383ace898 100644 --- a/mellea/formatters/granite/retrievers/util.py +++ b/mellea/formatters/granite/retrievers/util.py @@ -33,6 +33,7 @@ def download_mtrag_corpus(target_dir: str, corpus_name: str) -> pathlib.Path: Raises: ValueError: If `corpus_name` is not one of the supported corpus names. + urllib.error.HTTPError: On any HTTP error downloading the corpus file. """ corpus_names = ("cloud", "clapnq", "fiqa", "govt") if corpus_name not in corpus_names: @@ -44,7 +45,7 @@ def download_mtrag_corpus(target_dir: str, corpus_name: str) -> pathlib.Path: f"https://github.com/IBM/mt-rag-benchmark/raw/refs/heads/main/" f"corpora/{corpus_name}.jsonl.zip" ) - urllib.request.urlretrieve(source_url, target_file) + urllib.request.urlretrieve(source_url, str(target_file)) return target_file @@ -99,7 +100,9 @@ def read_mtrag_corpus(corpus_file: str | pathlib.Path) -> pa.Table: return t -def download_mtrag_embeddings(embedding_name: str, corpus_name: str, target_dir: str): +def download_mtrag_embeddings( + embedding_name: str, corpus_name: str, target_dir: str +) -> None: """Download precomputed embeddings for a corpus in the MTRAG benchmark. Args: @@ -114,6 +117,8 @@ def download_mtrag_embeddings(embedding_name: str, corpus_name: str, target_dir: ValueError: If `corpus_name` is not one of the supported corpus names, or if no precomputed embeddings are found for the given corpus and embedding model combination. + urllib.error.HTTPError: On any HTTP error other than 404 (which signals + end-of-parts). """ corpus_names = ("cloud", "clapnq", "fiqa", "govt") if corpus_name not in corpus_names: @@ -134,11 +139,13 @@ def download_mtrag_embeddings(embedding_name: str, corpus_name: str, target_dir: ) target_file = target_root / parquet_file_name try: - urllib.request.urlretrieve(source_url, target_file) + urllib.request.urlretrieve(source_url, str(target_file)) part_num += 1 - except urllib.error.HTTPError: - # Found all the parts; flow through - break + except urllib.error.HTTPError as exc: + if exc.code == 404: + # Found all the parts; flow through + break + raise # 429/5xx propagate rather than silently truncating if part_num == 1: raise ValueError( diff --git a/test/formatters/granite/test_retrievers_util.py b/test/formatters/granite/test_retrievers_util.py new file mode 100644 index 000000000..3633dde1e --- /dev/null +++ b/test/formatters/granite/test_retrievers_util.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for mellea.formatters.granite.retrievers.util download helpers.""" + +import http.client +import urllib.error +from unittest.mock import patch + +import pytest + +pytest.importorskip( + "pyarrow", reason="pyarrow not installed — install mellea[granite_retriever]" +) + +from mellea.formatters.granite.retrievers.util import ( # type: ignore[reportAttributeAccessIssue] + download_mtrag_corpus, + download_mtrag_embeddings, +) + + +def _http_error(code: int) -> urllib.error.HTTPError: + return urllib.error.HTTPError( + url="http://example.com", + code=code, + msg=str(code), + hdrs=http.client.HTTPMessage(), + fp=None, + ) + + +# --------------------------------------------------------------------------- +# download_mtrag_corpus +# --------------------------------------------------------------------------- + + +def test_corpus_raises_on_invalid_name(tmp_path): + with pytest.raises(ValueError, match="cloud"): + download_mtrag_corpus(str(tmp_path), "invalid_corpus") + + +def test_corpus_skips_download_if_file_exists(tmp_path): + target = tmp_path / "cloud.jsonl.zip" + target.write_bytes(b"dummy") + with patch("urllib.request.urlretrieve") as mock_retrieve: + download_mtrag_corpus(str(tmp_path), "cloud") + mock_retrieve.assert_not_called() + + +def test_corpus_downloads_when_file_missing(tmp_path): + with patch("urllib.request.urlretrieve") as mock_retrieve: + download_mtrag_corpus(str(tmp_path), "fiqa") + mock_retrieve.assert_called_once() + assert "fiqa" in mock_retrieve.call_args[0][0] + + +def test_corpus_propagates_http_error(tmp_path): + with patch("urllib.request.urlretrieve", side_effect=_http_error(429)): + with pytest.raises(urllib.error.HTTPError) as exc_info: + download_mtrag_corpus(str(tmp_path), "cloud") + assert exc_info.value.code == 429 + + +# --------------------------------------------------------------------------- +# download_mtrag_embeddings — 404 breaks cleanly, other errors propagate +# --------------------------------------------------------------------------- + + +def test_embeddings_404_on_first_part_raises_value_error(tmp_path): + """A 404 on part_001 means no embeddings exist — ValueError, not HTTPError.""" + with patch("urllib.request.urlretrieve", side_effect=_http_error(404)): + with pytest.raises(ValueError, match="No precomputed embeddings"): + download_mtrag_embeddings("model", "cloud", str(tmp_path)) + + +def test_embeddings_stops_after_404(tmp_path): + """Downloads part_001, then 404 on part_002 → returns cleanly with one part.""" + call_count = 0 + + def side_effect(_url, dest): + nonlocal call_count + call_count += 1 + if call_count == 1: + open(dest, "wb").close() # create the file + return + raise _http_error(404) + + with patch("urllib.request.urlretrieve", side_effect=side_effect): + download_mtrag_embeddings("model", "cloud", str(tmp_path)) + + assert call_count == 2 + + +@pytest.mark.parametrize("code", [429, 500, 502, 503, 504]) +def test_embeddings_propagates_non_404_errors(tmp_path, code): + """429/5xx must propagate — never silently treated as end-of-parts.""" + with patch("urllib.request.urlretrieve", side_effect=_http_error(code)): + with pytest.raises(urllib.error.HTTPError) as exc_info: + download_mtrag_embeddings("model", "cloud", str(tmp_path)) + assert exc_info.value.code == code