folding-studio/tests/query/test_protenix.py

import pytest
from folding_studio.query import ProtenixQuery


def test_from_file(tmp_files):
    """Test from_file with a valid FASTA file."""
    fasta_path = tmp_files["monomer_fasta"]
    query = ProtenixQuery.from_file(fasta_path, use_msa_server=True, seed=0)
    payload = query.payload
    assert fasta_path.stem in payload["fasta_files"]
    assert payload["use_msa_server"] is True
    assert payload["seeds"] == "0"


def test_from_file_invalid_extension(tmp_files):
    """Test from_file with an invalid file extension."""
    with pytest.raises(
        ValueError,
        match=f"Unsupported suffix '{tmp_files['invalid_source'].suffix}'",
    ):
        ProtenixQuery.from_file(
            tmp_files["invalid_source"], use_msa_server=True, seed=0
        )


def test_from_directory(tmp_files):
    """Test from_directory with valid FASTA files."""
    query = ProtenixQuery.from_directory(
        tmp_files["valid_dir"], use_msa_server=False, seed=0
    )
    payload = query.payload
    assert "monomer_1" in payload["fasta_files"]
    assert "monomer_2" in payload["fasta_files"]
    assert payload["use_msa_server"] is False
    assert payload["seeds"] == "0"


def test_from_empty_fasta_directory(tmp_files):
    """Test from_directory with an empty directory."""
    with pytest.raises(ValueError, match="No FASTA files found in directory"):
        ProtenixQuery.from_directory(
            tmp_files["empty_dir"], use_msa_server=False, seed=0
        )


def test_from_directory_with_invalid_sources(tmp_files):
    """Test from_directory ignores invalid file extensions."""
    with pytest.raises(ValueError, match="No FASTA files found in directory"):
        ProtenixQuery.from_directory(
            tmp_files["invalid_dir"], use_msa_server=False, seed=0
        )


def test_from_file_with_msa_server(tmp_files):
    """Test from_file with a FASTA file."""
    query = ProtenixQuery.from_file(
        str(tmp_files["monomer_fasta"]),
        use_msa_server=True,
        seed=10,
    )
    payload = query.payload
    assert tmp_files["monomer_fasta"].stem in payload["fasta_files"]
    assert payload["use_msa_server"] is True
    assert payload["seeds"] == "10"