folding-studio/tests/commands/predict/test_chai.py

from pathlib import Path

import pytest
from folding_studio.cli import app
from folding_studio.query import ChaiQuery
from typer.testing import CliRunner

runner = CliRunner()


RESTRAINTS_CSV_CONTENT = """
chainA,res_idxA,chainB,res_idxB,connection_type,confidence,min_distance_angstrom,max_distance_angstrom,comment,restraint_id
A,C387,B,Y101,contact,1.0,0.0,5.5,protein-heavy,restraint_1
C,I32,A,S483,contact,1.0,0.0,5.5,protein-light,restraint_2
"""


@pytest.fixture(scope="module")
def tmp_files(tmp_directory: Path, tmp_files: dict):
    """Generate temporary files."""

    # A valid restraints CSV file
    valid_restraints = tmp_directory / "example_restraints.csv"
    valid_restraints.write_text(RESTRAINTS_CSV_CONTENT)

    tmp_files["valid_restraints"] = valid_restraints

    yield tmp_files


def test_chai_fold_with_invalid_source(tmp_files):
    """
    Running the command with an unsupported file type should return an error.
    """
    result = runner.invoke(
        app,
        [
            "predict",
            "chai",
            str(tmp_files["invalid_source"]),
            "--project-code",
            "TEST_PROJECT",
        ],
    )
    assert result.exit_code != 0


def test_chai_fold_with_valid_fasta_file(
    mock_send_request, mock_download_results, tmp_files
):
    """
    Running the command with a valid FASTA file (without restraints) should
    process successfully and generate the correct payload.
    """
    result = runner.invoke(
        app,
        [
            "predict",
            "chai",
            str(tmp_files["monomer_fasta"]),
            "--project-code",
            "TEST_PROJECT",
            "--output",
            tmp_files["output_dir"],
        ],
    )
    assert result.exit_code == 0, result.output

    expected_query = ChaiQuery.from_file(
        path=str(tmp_files["monomer_fasta"]),
        use_msa_server=True,
        use_templates_server=False,
        num_trunk_recycles=3,
        seed=0,
        num_diffn_timesteps=200,
        restraints=None,
        recycle_msa_subsample=0,
        num_trunk_samples=1,
    )
    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 == expected

    mock_download_results.assert_called_once()


def test_chai_fold_with_valid_fasta_and_restraints(
    mock_send_request, mock_download_results, tmp_files
):
    """
    Running the command with a valid FASTA file and a valid restraints CSV file
    should include the restraints content in the payload.
    """
    result = runner.invoke(
        app,
        [
            "predict",
            "chai",
            str(tmp_files["monomer_fasta"]),
            "--project-code",
            "TEST_PROJECT",
            "--restraints",
            str(tmp_files["valid_restraints"]),
            "--output",
            tmp_files["output_dir"],
        ],
    )
    assert result.exit_code == 0, result.output

    restraints_content = tmp_files["valid_restraints"].read_text().strip()

    expected_query = ChaiQuery.from_file(
        path=str(tmp_files["monomer_fasta"]),
        use_msa_server=True,
        use_templates_server=False,
        num_trunk_recycles=3,
        seed=0,
        num_diffn_timesteps=200,
        restraints=str(tmp_files["valid_restraints"]),
        recycle_msa_subsample=0,
        num_trunk_samples=1,
    )
    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 == expected
    assert payload["restraints"] == restraints_content

    mock_download_results.assert_called_once()


def test_chai_fold_with_valid_fasta_directory(
    mock_send_request, mock_download_results, tmp_files
):
    """
    Running the command with a directory containing valid FASTA files should
    process successfully and submit a single request for the directory.
    """
    result = runner.invoke(
        app,
        [
            "predict",
            "chai",
            str(tmp_files["valid_dir"]),
            "--project-code",
            "TEST_PROJECT",
            "--output",
            tmp_files["output_dir"],
        ],
    )
    assert result.exit_code == 0, result.output

    expected_query = ChaiQuery.from_directory(
        path=str(tmp_files["valid_dir"]),
        use_msa_server=True,
        use_templates_server=False,
        num_trunk_recycles=3,
        seed=0,
        num_diffn_timesteps=200,
        restraints=None,
        recycle_msa_subsample=0,
        num_trunk_samples=1,
    )

    expected_fasta_files = expected_query.payload["fasta_files"]

    actual_fasta_files = [
        call_args[0][0].payload["fasta_files"]
        for call_args in mock_send_request.call_args_list
    ]

    # Check that the expected FASTA files match one of the actual request payloads
    assert expected_fasta_files in actual_fasta_files, (
        f"FASTA file {expected_fasta_files} was expected but not found in actual requests."
    )

    mock_download_results.assert_called()



def test_chai_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",
            "chai",
            str(tmp_files["empty_dir"]),
            "--project-code",
            "TEST_PROJECT",
            "--output",
            tmp_files["output_dir"],
        ],
    )
    assert result.exit_code != 0, result.output