125 lines
4.0 KiB
Python
Executable File
125 lines
4.0 KiB
Python
Executable File
from __future__ import annotations
|
|
|
|
import http.client
|
|
import zipfile
|
|
|
|
|
|
def test_reference_bundle_download_reports_progress_events(tmp_path, monkeypatch):
|
|
from apps.web import pipeline_backend as backend
|
|
|
|
job = backend.JobInfo(
|
|
job_id="job_test_ref_bundle_progress",
|
|
input_dir=str(tmp_path / "input"),
|
|
workspace_dir=str(tmp_path / "workspace"),
|
|
logs_dir=str(tmp_path / "workspace" / "logs"),
|
|
created_at="",
|
|
)
|
|
input_dir = tmp_path / "input"
|
|
workspace = tmp_path / "workspace"
|
|
(workspace / "final").mkdir(parents=True)
|
|
(workspace / "logs").mkdir(parents=True)
|
|
input_dir.mkdir(parents=True)
|
|
|
|
fake_zip = tmp_path / "ncbi_dataset.zip"
|
|
with zipfile.ZipFile(fake_zip, "w") as zf:
|
|
zf.writestr("ncbi_dataset/data/GCF_FAKE/protein.faa", ">SOE_PEP1\nMAAA\n")
|
|
zf.writestr(
|
|
"ncbi_dataset/data/GCF_FAKE/cds_from_genomic.fna",
|
|
">lcl|NC_000001.1_cds_XP_000001.1_1 [protein_id=XP_000001.1]\nATGGCG\n",
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
backend,
|
|
"_select_ncbi_reference_assembly",
|
|
lambda **kwargs: {
|
|
"accession": "GCF_FAKE",
|
|
"organism_name": "Spinacia oleracea",
|
|
"taxon_id": "3562",
|
|
"assembly_name": "fake",
|
|
"assembly_level": "Chromosome",
|
|
"annotation_name": "fake annotation",
|
|
},
|
|
)
|
|
|
|
def fake_download(_url, output_path, timeout_sec=300, progress_callback=None):
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
data = fake_zip.read_bytes()
|
|
if progress_callback:
|
|
progress_callback({
|
|
"phase": "download",
|
|
"bytes_downloaded": len(data),
|
|
"total_bytes": len(data),
|
|
"percentage": 100,
|
|
})
|
|
output_path.write_bytes(data)
|
|
return output_path.stat().st_size
|
|
|
|
monkeypatch.setattr(backend, "_download_url_to_file", fake_download)
|
|
events = []
|
|
|
|
manifest = backend.download_reference_annotation_bundle(
|
|
job=job,
|
|
species_code="soe",
|
|
species_name="Spinacia oleracea",
|
|
taxon_id="3562",
|
|
requested_stages=["ppi"],
|
|
progress_callback=events.append,
|
|
)
|
|
|
|
phases = [event.get("phase") for event in events]
|
|
assert manifest["status"] == "success"
|
|
assert "select_assembly" in phases
|
|
assert "download" in phases
|
|
assert "extract" in phases
|
|
assert "validate_ids" in phases
|
|
assert phases[-1] == "completed"
|
|
assert manifest["requested_stages"] == ["ppi"]
|
|
assert manifest["requested_file_types"] == ["pep", "cds"]
|
|
assert "cds" in manifest["files"]
|
|
|
|
|
|
def test_reference_bundle_download_retries_incomplete_read(tmp_path, monkeypatch):
|
|
from apps.web import pipeline_backend as backend
|
|
|
|
payload = b"complete archive"
|
|
calls = {"count": 0}
|
|
|
|
class FakeResponse:
|
|
def __init__(self, *, incomplete: bool):
|
|
self.incomplete = incomplete
|
|
self.headers = {"Content-Length": str(len(payload))}
|
|
self._sent = False
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *_args):
|
|
return False
|
|
|
|
def read(self, _size):
|
|
if self.incomplete:
|
|
raise http.client.IncompleteRead(b"partial", len(payload))
|
|
if self._sent:
|
|
return b""
|
|
self._sent = True
|
|
return payload
|
|
|
|
def fake_urlopen(_url, timeout):
|
|
calls["count"] += 1
|
|
return FakeResponse(incomplete=calls["count"] == 1)
|
|
|
|
monkeypatch.setattr(backend.urllib.request, "urlopen", fake_urlopen)
|
|
monkeypatch.setattr(backend.time, "sleep", lambda _seconds: None)
|
|
output_path = tmp_path / "reference.zip"
|
|
|
|
downloaded = backend._download_url_to_file(
|
|
"https://example.invalid/reference.zip",
|
|
output_path,
|
|
attempts=2,
|
|
)
|
|
|
|
assert calls["count"] == 2
|
|
assert downloaded == len(payload)
|
|
assert output_path.read_bytes() == payload
|
|
assert not output_path.with_suffix(".zip.part").exists()
|