File size: 3,285 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 |
"""Test upload custom files."""
import datetime
from unittest import mock
from unittest.mock import patch
import pytest
from pathlib import Path
import requests_mock
from folding_studio_data_models import CustomFileType
from folding_studio.api_call.upload_custom_files import upload_custom_files
from folding_studio.config import API_URL
fixed_date = datetime.datetime(2024, 11, 18, 12, 34, 56, 123456)
formatted_fixed_date = fixed_date.strftime("%Y%m%d%H%M%S%f")
# Helper to mock the signed URL request
def mock_signed_url_request(
mocker, filename: str, signed_url: str, file_type: str, bucket: str
):
"""Helper to mock signed URL API response."""
file_base, ext = str(filename).split(".")
url = f"{API_URL}getUploadSignedURL?blob_name={filename}&file_type={file_type}"
return mocker.get(
url=url,
json={
"signed_url": signed_url,
"destination_file": f"gs://{bucket}/{file_base}_d41d8cd9.{ext}",
"destination_bucket": f"gs://{bucket}",
},
)
# Parameterized test to handle different file types and buckets
@pytest.mark.parametrize(
"filenames, bucket, file_type_enum",
[
(
["file_1.cif", "file_2.cif", "file_3.cif"],
"custom_templates",
CustomFileType.TEMPLATE,
),
(
["file_1.sto", "file_2.a3m"],
"custom_msas",
CustomFileType.MSA,
),
(
["file_1.cif", "file_2.cif"],
"initial_guess",
CustomFileType.INITIAL_GUESS,
),
(
["file_1.json", "file_2.json"],
"templates_masks",
CustomFileType.TEMPLATE_MASK,
),
(
["file_1.aligned.pqt", "file_2.aligned.pqt"],
"custom_msas",
CustomFileType.MSA
),
],
)
def test_upload_custom_files_pass(
tmp_path: Path,
filenames: list,
bucket: str,
file_type_enum: CustomFileType,
headers: dict[str, str]
):
"""Test upload custom files pass for different file types and buckets."""
token = "identity_token"
file_paths = [tmp_path / filename for filename in filenames]
# Create the files
for file in file_paths:
file.touch()
signed_url = "https://storage.googleapis.com/some_url"
with patch(
"folding_studio.api_call.upload_custom_files._get_blob_zip_name"
) as mock_blob_name, requests_mock.Mocker() as m:
# Mock the signed URL requests for each file
for _ in filenames:
mock_blob_name.return_value = "mock_file.zip"
mock_signed_url_request(
mocker=m,
filename="mock_file.zip",
signed_url=signed_url,
file_type=file_type_enum.value,
bucket=bucket,
)
# Mock the PUT request to upload the file
m.put(signed_url, status_code=200)
m.post(API_URL + "unzipFileInBucket", status_code=200)
# Call the function being tested
output = upload_custom_files(headers, file_paths, file_type=file_type_enum)
# Assertions
assert set(output.keys()) == {str(f) for f in file_paths}
assert all(f"gs://{bucket}/" in val for val in output.values())
|