From 19e87e86ee843351f0a337e8a3e818f858c4d716 Mon Sep 17 00:00:00 2001 From: laclac102 <53896516+vtnphan@users.noreply.github.com> Date: Fri, 25 Sep 2026 09:39:31 +1000 Subject: [PATCH 1/2] fix: collect confidence json files for wisps --- app/services/results_utils.py | 10 ++- tests/test_services_results_utils.py | 94 +++++++++++++++++++++++++++- 2 files changed, 101 insertions(+), 3 deletions(-) diff --git a/app/services/results_utils.py b/app/services/results_utils.py index 5a155c8a..daf31973 100644 --- a/app/services/results_utils.py +++ b/app/services/results_utils.py @@ -328,7 +328,7 @@ async def resolve_fasta_form_data( response_content_disposition=_format_attachment_content_disposition(filename), settings=settings, ) - except S3ConfigurationError, S3ServiceError: + except (S3ConfigurationError, S3ServiceError): logger.warning( "Failed to generate presigned URL for %r (S3 key %r)", key, @@ -371,7 +371,7 @@ async def resolve_pdb_presigned_urls( settings=settings, ) return {**form_data, "starting_pdb": presigned_url} - except S3ConfigurationError, S3ServiceError: + except (S3ConfigurationError, S3ServiceError): logger.warning( "Failed to generate presigned starting_pdb URL for S3 key %r; " "returning original form data", @@ -789,6 +789,11 @@ 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 +870,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..3cf07676 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,39 @@ 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 +1660,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 +1669,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 +1786,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 From db3e872d6a22dd82329dfbe451c615db94c67be2 Mon Sep 17 00:00:00 2001 From: laclac102 <53896516+vtnphan@users.noreply.github.com> Date: Fri, 25 Sep 2026 10:00:32 +1000 Subject: [PATCH 2/2] chore: lint --- app/services/results_utils.py | 8 +++----- tests/test_services_results_utils.py | 4 +--- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/app/services/results_utils.py b/app/services/results_utils.py index daf31973..7c3ad63b 100644 --- a/app/services/results_utils.py +++ b/app/services/results_utils.py @@ -328,7 +328,7 @@ async def resolve_fasta_form_data( response_content_disposition=_format_attachment_content_disposition(filename), settings=settings, ) - except (S3ConfigurationError, S3ServiceError): + except S3ConfigurationError, S3ServiceError: logger.warning( "Failed to generate presigned URL for %r (S3 key %r)", key, @@ -371,7 +371,7 @@ async def resolve_pdb_presigned_urls( settings=settings, ) return {**form_data, "starting_pdb": presigned_url} - except (S3ConfigurationError, S3ServiceError): + except S3ConfigurationError, S3ServiceError: logger.warning( "Failed to generate presigned starting_pdb URL for S3 key %r; " "returning original form data", @@ -789,9 +789,7 @@ 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" - ): + 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 diff --git a/tests/test_services_results_utils.py b/tests/test_services_results_utils.py index 3cf07676..11304026 100644 --- a/tests/test_services_results_utils.py +++ b/tests/test_services_results_utils.py @@ -1489,9 +1489,7 @@ def test_classify_wisps_output_key_colabfold_confidence_json_keeps_existing_name None, "colabfold", ) - assert result == ClassifiedOutput( - category="pae", label="sample1_model_0_confidence.json" - ) + assert result == ClassifiedOutput(category="pae", label="sample1_model_0_confidence.json") def test_classify_wisps_output_key_boltz_confidence_json_not_classified():