|
|
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 |
|
|
|
|
|
|
|
|
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] |
|
|
|
|
|
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 |
|
|
|
