diff --git a/app/services/results_utils.py b/app/services/results_utils.py index 5a155c8a..7c3ad63b 100644 --- a/app/services/results_utils.py +++ b/app/services/results_utils.py @@ -789,6 +789,9 @@ def classify_wisps_output_key( return ClassifiedOutput(category="pdb", label=basename) if "/colabfold_predictions/pae/" in lowered and basename.lower().endswith(".npz"): return ClassifiedOutput(category="pae", label=basename) + if "/colabfold_predictions/confidence/" in lowered and basename.lower().endswith(".json"): + label = basename if "confidence" in basename.lower() else f"confidence_{basename}" + return ClassifiedOutput(category="pae", label=label) return None @@ -865,6 +868,7 @@ def build_wisps_output_listing_prefixes(run: WorkflowRun) -> list[str]: if tool in ("", "colabfold"): prefixes.append(f"{run_uuid}/colabfold_predictions/pdb/") prefixes.append(f"{run_uuid}/colabfold_predictions/pae/") + prefixes.append(f"{run_uuid}/colabfold_predictions/confidence/") return prefixes diff --git a/tests/test_services_results_utils.py b/tests/test_services_results_utils.py index 96d89ac1..11304026 100644 --- a/tests/test_services_results_utils.py +++ b/tests/test_services_results_utils.py @@ -315,6 +315,7 @@ def test_builtin_specs_get_transfer_prefixes_excludes_run_root_but_keeps_usage_r f"{run.id}/boltz_predictions/pae/", f"{run.id}/colabfold_predictions/pdb/", f"{run.id}/colabfold_predictions/pae/", + f"{run.id}/colabfold_predictions/confidence/", f"{run.id}/UsageReport.csv", ] assert rfdiffusion_spec.get_transfer_prefixes(run) == [ @@ -1470,6 +1471,37 @@ def test_classify_wisps_output_key_colabfold_pae_npz(): assert result == ClassifiedOutput(category="pae", label="sample1_model_0.npz") +def test_classify_wisps_output_key_colabfold_confidence_json(): + """Confidence JSONs join the existing (hidden) "pae" bucket rather than a + new category, and get a "confidence_" prefix when the filename doesn't + already say so, so they're distinguishable from the .npz PAE files.""" + run_id = str(uuid4()) + result = classify_wisps_output_key( + f"{run_id}/colabfold_predictions/confidence/sample1_model_0.json", None, "colabfold" + ) + assert result == ClassifiedOutput(category="pae", label="confidence_sample1_model_0.json") + + +def test_classify_wisps_output_key_colabfold_confidence_json_keeps_existing_name(): + run_id = str(uuid4()) + result = classify_wisps_output_key( + f"{run_id}/colabfold_predictions/confidence/sample1_model_0_confidence.json", + None, + "colabfold", + ) + assert result == ClassifiedOutput(category="pae", label="sample1_model_0_confidence.json") + + +def test_classify_wisps_output_key_boltz_confidence_json_not_classified(): + """Confidence JSON collection is colabfold-specific; boltz runs don't + produce this folder, so a stray file there must not be picked up.""" + run_id = str(uuid4()) + result = classify_wisps_output_key( + f"{run_id}/colabfold_predictions/confidence/sample1_model_0.json", None, "boltz" + ) + assert result is None + + def test_classify_wisps_output_key_tool_restricts_to_matching_predictions_folder(): """A run's classifier only matches its own tool's prediction folder - the other tool's folder never exists for that run, so it must never be misclassified.""" @@ -1626,7 +1658,7 @@ async def test_get_run_service_usage(): def test_build_wisps_output_listing_prefixes(): run = WorkflowRun(id=uuid4(), owner_user_id=uuid4(), sample_id="sample1") prefixes = build_wisps_output_listing_prefixes(run) - assert len(prefixes) == 8 + assert len(prefixes) == 9 assert f"{run.id}/" in prefixes assert f"{run.id}/multiqc/" in prefixes assert f"{run.id}/collect/" in prefixes @@ -1635,6 +1667,17 @@ def test_build_wisps_output_listing_prefixes(): assert f"{run.id}/colabfold_predictions/pdb/" in prefixes assert f"{run.id}/boltz_predictions/pae/" in prefixes assert f"{run.id}/colabfold_predictions/pae/" in prefixes + assert f"{run.id}/colabfold_predictions/confidence/" in prefixes + + +def test_build_wisps_output_listing_prefixes_colabfold_tool_only(): + run = SimpleNamespace(id=uuid4(), tool="colabfold", submitted_form_data=None) + prefixes = build_wisps_output_listing_prefixes(run) + assert f"{run.id}/colabfold_predictions/pdb/" in prefixes + assert f"{run.id}/colabfold_predictions/pae/" in prefixes + assert f"{run.id}/colabfold_predictions/confidence/" in prefixes + assert f"{run.id}/boltz_predictions/cif/" not in prefixes + assert f"{run.id}/boltz_predictions/pae/" not in prefixes def test_build_wisps_output_listing_prefixes_returns_empty_when_no_id(): @@ -1741,6 +1784,53 @@ async def read_bytes(key: str, **_kwargs) -> bytes: assert zip_file.read(f"pdb/{structure_filename}") == structure_content +@pytest.mark.asyncio +async def test_get_category_downloads_zipped_bundles_colabfold_confidence_with_pae( + test_db, persistent_models +): + """Colabfold confidence/*.json files join the existing "pae" bucket, so + the "pae" bundle download for a colabfold run contains both the .npz + files and the confidence JSONs (renamed to make them distinguishable).""" + user = AppUserFactory.create_sync() + run = WorkflowRunFactory.create_sync( + owner=user, + workflow=Workflow( + name="interaction-screening", + repo_url="https://github.com/test/interaction-screening", + default_revision="main", + config_path="/config/interaction-screening.config", + ), + tool="colabfold", + seqera_run_id="wf-wisps-zip-confidence", + ) + output_contents = { + f"{run.id}/multiqc/multiqc_report.html": b"report", + f"{run.id}/colabfold_predictions/pae/sample1_model_0.npz": b"npz-bytes", + f"{run.id}/colabfold_predictions/confidence/sample1_model_0.json": b'{"ptm": 0.9}', + } + outputs = [S3Object(object_key=key, uri=f"s3://bucket/{key}") for key in output_contents] + test_db.add_all([user, run, *outputs]) + test_db.commit() + test_db.add_all([_make_run_output(run, item.object_key) for item in outputs]) + test_db.commit() + + async def read_bytes(key: str, **_kwargs) -> bytes: + return output_contents[key] + + with patch( + "app.services.results_utils.read_s3_bytes", + new=AsyncMock(side_effect=read_bytes), + ): + zip_buffer = await get_category_downloads_zipped(test_db, run, "pae") + + with ZipFile(BytesIO(zip_buffer.getvalue())) as zip_file: + assert set(zip_file.namelist()) == { + "pae/sample1_model_0.npz", + "pae/confidence_sample1_model_0.json", + } + assert zip_file.read("pae/confidence_sample1_model_0.json") == b'{"ptm": 0.9}' + + @pytest.mark.asyncio async def test_get_category_downloads_zipped_returns_none_without_matches( test_db, persistent_models