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

from pathlib import Path
from unittest import mock

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

current_workdir = Path(__file__).parent.resolve()
data_dir = Path(current_workdir / "data")

runner = CliRunner()


def test_predict_with_fasta_file_pass(
    mock_send_request: mock.Mock,
    mock_download_results: mock.Mock,
    tmp_files,
):
    result = runner.invoke(
        app,
        [
            "predict",
            "soloseq",
            str(tmp_files["monomer_fasta"]),
            "--output",
            tmp_files["output_dir"],
            "--project-code",
            "my_project_code",
        ],
    )
    assert result.exit_code == 0, result.output

    expected_query = SoloSeqQuery(
        fasta_files={"monomer": ">tag1|tag2\nABCDEGF"}, query_name="monomer"
    )
    mock_send_request.assert_called_once_with(expected_query, "my_project_code")
    mock_download_results.assert_called_once()


def test_predict_with_fasta_dir_pass(
    mock_send_request: mock.Mock,
    mock_download_results: mock.Mock,
    tmp_files,
):
    result = runner.invoke(
        app,
        [
            "predict",
            "soloseq",
            str(tmp_files["valid_dir"]),
            "--output",
            tmp_files["output_dir"],
            "--project-code",
            "my_project_code",
        ],
    )
    assert result.exit_code == 0, result.output

    expected_query = SoloSeqQuery(
        fasta_files={
            "monomer_1": ">tag1|tag2\nABCDEGF",
            "monomer_2": ">tag1|tag2\nABCDEGF",
        },
        query_name="valid_dir",
    )
    mock_send_request.assert_called_once_with(expected_query, "my_project_code")
    mock_download_results.assert_called_once()