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

from folding_studio.cli import app
from folding_studio.query import ProtenixQuery
from typer.testing import CliRunner

runner = CliRunner()


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


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

    # Create the expected query using the same parameters as defaults.
    expected_query = ProtenixQuery.from_file(tmp_files["monomer_fasta"], seed=42)
    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
    assert payload["fasta_files"] == expected["fasta_files"]
    assert payload["use_msa_server"] == expected["use_msa_server"]
    assert payload["seeds"] == expected["seeds"]

    mock_download_results.assert_called_once()


def test_protenix_predict_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 combine all FASTA files into the payload.
    """
    result = runner.invoke(
        app,
        [
            "predict",
            "protenix",
            str(tmp_files["valid_dir"]),
            "--project-code",
            "TEST_PROJECT",
            "--output",
            tmp_files["output_dir"],
        ],
    )
    assert result.exit_code == 0, result.output

    expected_query = ProtenixQuery.from_directory(tmp_files["valid_dir"])
    mock_send_request.assert_called_once()
    actual_query = mock_send_request.call_args[0][0]
    payload = actual_query.payload
    expected = expected_query.payload
    # For directories, the fasta_files dictionary should contain an entry for each FASTA file.
    assert payload["fasta_files"] == expected["fasta_files"]

    mock_download_results.assert_called_once()


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