folding-studio/tests/conftest.py

import os
import shutil
from pathlib import Path
from unittest import mock


import pytest
import yaml
from folding_studio_data_models.request.folding import FoldingModel


@pytest.fixture(scope="session", autouse=True)
def set_env():
    """
    We set the env var here instead of pytest.ini, because
    we want to override the env vars in some tests.
    I could not find an easy way to override the content of pytest.ini.
    """
    os.environ["API_URL"] = "https://test_api_url/"
    os.environ["FOLDING_PROJECT_CODE"] = "FOLDING_DEV"
    os.environ["FOLDING_API_KEY"] = "MY_KEY"


@pytest.fixture
def remove_project_code_from_env_var(monkeypatch):
    with mock.patch.dict(os.environ):
        monkeypatch.delenv("FOLDING_PROJECT_CODE")
        yield

@pytest.fixture
def remove_api_key_from_env_var(monkeypatch):
    with mock.patch.dict(os.environ):
        monkeypatch.delenv("FOLDING_API_KEY")
        yield

@pytest.fixture(params=[FoldingModel.AF2, FoldingModel.OPENFOLD])
def folding_model(request):
    return request.param

@pytest.fixture
def headers():
    return {"Authorization": "Bearer identity_token"}


@pytest.fixture(scope="module")
def tmp_directory(tmp_path_factory):
    tmp_dir_path = tmp_path_factory.mktemp("data")
    yield tmp_dir_path
    shutil.rmtree(tmp_dir_path)


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

    # generate fasta files
    empty_fasta = tmp_directory / "empty.fasta"
    empty_fasta.touch()

    monomer_fasta = tmp_directory / "monomer.fasta"
    with monomer_fasta.open("w") as f:
        f.write(">tag1|tag2\nABCDEGF")

    multimer_fasta = tmp_directory / "multimer.fasta"
    with multimer_fasta.open("w") as f:
        f.write(">tag1\nABCDEGF\n>tag2\nABCDEGF")

    # generate yaml file
    yaml_content = {
        "version": 1,
        "sequences": [{"protein": {"id": "A", "sequence": "QLEDSEVEAVAKGLEE"}}],
    }
    yaml_file_path = tmp_directory / "protein.yaml"
    with yaml_file_path.open("w") as f:
        yaml.safe_dump(yaml_content, f, default_flow_style=False)

    # invalid source
    invalid_source = tmp_directory / "protein.txt"
    invalid_source.touch()

    # batch directory with invalid source files
    dir_with_multimer = tmp_directory / "dir_with_multimer"
    dir_with_multimer.mkdir(parents=True, exist_ok=True)
    shutil.copy(multimer_fasta, dir_with_multimer / "multimer.fasta")
    shutil.copy(monomer_fasta, dir_with_multimer / "monomer.fasta")

    invalid_dir = tmp_directory / "invalid_dir"
    invalid_dir.mkdir(parents=True, exist_ok=True)
    (invalid_dir / "protein_A.txt").touch()
    (invalid_dir / "protein_B.txt").touch()

    mixed_fasta_txt_dir = tmp_directory / "mixed_fasta_txt_dir"
    mixed_fasta_txt_dir.mkdir(parents=True, exist_ok=True)
    (mixed_fasta_txt_dir / "protein_A.fasta").touch()
    (mixed_fasta_txt_dir / "protein_B.txt").touch()

    mixed_fasta_yaml_dir = tmp_directory / "mixed_fasta_yaml_dir"
    mixed_fasta_yaml_dir.mkdir(parents=True, exist_ok=True)
    shutil.copy(monomer_fasta, mixed_fasta_yaml_dir / "protein_A.fasta")
    shutil.copy(yaml_file_path, mixed_fasta_yaml_dir / "protein_B.yaml")

    empty_dir = tmp_directory / "empty_dir"
    empty_dir.mkdir(parents=True, exist_ok=True)

    # valid batch directory
    valid_dir = tmp_directory / "valid_dir"
    valid_dir.mkdir(parents=True, exist_ok=True)
    shutil.copy(monomer_fasta, valid_dir / "monomer_1.fasta")
    shutil.copy(monomer_fasta, valid_dir / "monomer_2.fasta")

    files = {
        "invalid_source": invalid_source,
        "empty_fasta": empty_fasta,
        "monomer_fasta": monomer_fasta,
        "multimer_fasta": multimer_fasta,
        "yaml_file_path": yaml_file_path,
        "valid_dir": valid_dir,
        "invalid_dir": invalid_dir,
        "dir_with_multimer": dir_with_multimer,
        "mixed_dir": mixed_fasta_txt_dir,
        "empty_dir": empty_dir,
        "output_dir": tmp_directory / "output",
        "mixed_fasta_yaml_dir": mixed_fasta_yaml_dir,
    }

    yield files