comet-ml--opik/sdks/python/tests/unit/file_upload/conftest.py

import tempfile
from unittest import mock

import numpy as np
import pytest
import uuid6

from opik.file_upload import (
    upload_client,
)
from opik.message_processing import messages
from opik.rest_api import types as rest_api_types

FILE_SIZE = 12 * 1024 * 1024


@pytest.fixture
def data_file():
    with tempfile.NamedTemporaryFile(delete=True) as file:
        file.write(np.random.bytes(FILE_SIZE))
        file.seek(0)

        yield file


@pytest.fixture
def attachment(data_file):
    attachment = messages.CreateAttachmentMessage(
        file_path=data_file.name,
        file_name="test_attachment.dat",
        mime_type="application/octet-stream",
        entity_type="span",
        entity_id="entity_id",
        project_name="project_name",
        encoded_url_override="encoded_url_override_path",
    )
    yield attachment


@pytest.fixture
def permanent_data_file():
    with tempfile.NamedTemporaryFile(delete=False) as file:
        file.write(np.random.bytes(FILE_SIZE))
        file.seek(0)

        yield file


@pytest.fixture
def attachment_to_be_deleted(permanent_data_file):
    attachment = messages.CreateAttachmentMessage(
        file_path=permanent_data_file.name,
        file_name="test_attachment.dat",
        mime_type="application/octet-stream",
        entity_type="span",
        entity_id="entity_id",
        project_name="project_name",
        encoded_url_override="encoded_url_override_path",
        delete_after_upload=True,
    )
    yield attachment


@pytest.fixture
def rest_client_s3():
    # returns response to start_upload request initiating S3 upload
    pre_sign_urls = [
        "https://s3.amazonaws.com/bucket/1",
        "https://s3.amazonaws.com/bucket/2",
        "https://s3.amazonaws.com/bucket/3",
    ]
    start_upload_response = rest_api_types.StartMultipartUploadResponse(
        upload_id=str(uuid6.uuid7()), pre_sign_urls=pre_sign_urls
    )
    rest_client = mock.Mock()
    rest_client.attachments.start_multi_part_upload.return_value = start_upload_response

    yield rest_client


@pytest.fixture
def rest_client_local():
    pre_sign_urls = [
        "https://localhost:8080/bucket/1",
    ]
    start_upload_response = rest_api_types.StartMultipartUploadResponse(
        upload_id=upload_client.LOCAL_UPLOAD_MAGIC_ID, pre_sign_urls=pre_sign_urls
    )
    rest_client = mock.Mock()
    rest_client.attachments.start_multi_part_upload.return_value = start_upload_response

    yield rest_client