File size: 3,689 Bytes
44459bb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
from unittest import mock
import pytest
from folding_studio.cli import app
from folding_studio.query import BoltzQuery
from typer.testing import CliRunner
runner = CliRunner()
@pytest.fixture
def inference_parameters():
"""Provides the inference parameters dictionary to tests."""
return {
"recycling_steps": 3,
"sampling_steps": 200,
"diffusion_samples": 1,
"step_scale": 1.638,
"output_format": "mmcif",
"num_workers": 2,
"msa_pairing_strategy": "greedy",
"write_full_pae": False,
"write_full_pde": False,
"seed": 0,
}
@pytest.fixture()
def mock_save_inference_parameters():
"""Monkeypatch Response.download_results so that no real file download occurs."""
with mock.patch("folding_studio.query.Query.save_parameters") as m:
yield m
def test_boltz_fold_with_valid_fasta_file(
mock_send_request,
mock_download_results,
mock_save_inference_parameters,
tmp_files,
inference_parameters,
):
"""
Running the command with a valid FASTA file should
process successfully and generate the correct payload.
"""
result = runner.invoke(
app,
[
"predict",
"boltz",
str(tmp_files["monomer_fasta"]),
"--project-code",
"test_project",
"--output",
tmp_files["output_dir"],
],
)
assert result.exit_code == 0, result.output
# Create the expected query using the same parameters as defaults.
expected_query = BoltzQuery.from_file(
path=str(tmp_files["monomer_fasta"]), **inference_parameters
)
mock_send_request.assert_called_once()
actual_query = mock_send_request.call_args[0][0]
# Compare payload keys
payload = actual_query.payload
expected = expected_query.payload
for key in ["fasta_files", "yaml_files", "parameters"]:
assert payload[key] == expected[key]
mock_save_inference_parameters.assert_called_once()
mock_download_results.assert_called_once()
def test_boltz_fold_with_valid_fasta_directory(
mock_send_request,
mock_download_results,
mock_save_inference_parameters,
tmp_files,
inference_parameters,
):
"""
Running the command with a directory containing valid FASTA files should
process successfully and combine all FASTA files into the payload.
"""
result = runner.invoke(
app,
[
"predict",
"boltz",
str(tmp_files["valid_dir"]),
"--project-code",
"test_project",
"--output",
tmp_files["output_dir"],
],
)
assert result.exit_code == 0, result.output
expected_query = BoltzQuery.from_directory(
path=tmp_files["valid_dir"], **inference_parameters
)
mock_send_request.assert_called_once()
actual_query = mock_send_request.call_args[0][0]
payload = actual_query.payload
expected = expected_query.payload
assert payload["fasta_files"] == expected["fasta_files"]
mock_save_inference_parameters.assert_called_once()
mock_download_results.assert_called_once()
def test_boltz_fold_with_empty_directory(tmp_files):
"""
Running the command with an empty directory should return an error,
since no FASTA files are available.
"""
result = runner.invoke(
app,
[
"predict",
"boltz",
str(tmp_files["empty_dir"]),
"--project-code",
"test_project",
"--output",
tmp_files["output_dir"],
],
)
assert result.exit_code != 0, result.output
|