danswer/backend/tests/integration/common_utils/managers/file.py

import mimetypes
from typing import cast
from typing import IO
from typing import List
from typing import Tuple

import requests

from onyx.file_store.models import FileDescriptor
from tests.integration.common_utils.constants import API_SERVER_URL
from tests.integration.common_utils.constants import GENERAL_HEADERS
from tests.integration.common_utils.test_models import DATestUser


class FileManager:
    @staticmethod
    def upload_files(
        files: List[Tuple[str, IO]],
        user_performing_action: DATestUser | None = None,
    ) -> Tuple[List[FileDescriptor], str]:
        headers = (
            user_performing_action.headers
            if user_performing_action
            else GENERAL_HEADERS
        )
        headers.pop("Content-Type", None)

        files_param = []
        for filename, file_obj in files:
            mime_type, _ = mimetypes.guess_type(filename)
            if mime_type is None:
                mime_type = "application/octet-stream"
            files_param.append(("files", (filename, file_obj, mime_type)))

        response = requests.post(
            f"{API_SERVER_URL}/chat/file",
            files=files_param,
            headers=headers,
        )

        if not response.ok:
            return (
                cast(List[FileDescriptor], []),
                f"Failed to upload files - {response.json().get('detail', 'Unknown error')}",
            )

        response_json = response.json()
        return response_json.get("files", cast(List[FileDescriptor], [])), ""

    @staticmethod
    def fetch_uploaded_file(
        file_id: str,
        user_performing_action: DATestUser | None = None,
    ) -> bytes:
        response = requests.get(
            f"{API_SERVER_URL}/chat/file/{file_id}",
            headers=user_performing_action.headers
            if user_performing_action
            else GENERAL_HEADERS,
        )
        response.raise_for_status()
        return response.content