folding-studio/tests/api_call/test_upload_custom_files.py

"""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())