File size: 4,065 Bytes
44459bb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
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
|