comet-ml--opik/sdks/python/tests/unit/api_objects/attachment/test_converters.py

import os
import tempfile
from typing import Generator
from unittest import mock

import pytest

from opik.api_objects import attachment
from opik.api_objects.attachment import converters
from opik.message_processing import messages


@pytest.fixture
def original_file() -> Generator[str, None, None]:
    """Create a file with test content and clean up after test."""
    with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
        f.write("test content")
        path = f.name

    yield path

    if os.path.exists(path):
        os.unlink(path)


@pytest.mark.parametrize(
    "attachment_data,expected",
    [
        (attachment.Attachment(data="test.png", file_name=None), "image/png"),
        (attachment.Attachment(data="test.png", file_name="test.jpg"), "image/jpeg"),
        (
            attachment.Attachment(
                data="test.pdf", file_name=None, content_type="image/jpeg"
            ),
            "image/jpeg",
        ),
    ],
)
def test_guess_attachment_type(attachment_data: attachment.Attachment, expected: str):
    mimetype = converters.guess_attachment_type(attachment_data)
    assert mimetype == expected


def test_attachment_to_message__no_temp_copy(original_file: str):
    url_override = "https://example.com"
    entity_id = "123"
    project_name = "test-project"
    attachment_data = attachment.Attachment(data=original_file, create_temp_copy=False)

    message = converters.attachment_to_message(
        attachment_data=attachment_data,
        entity_type="trace",
        entity_id=entity_id,
        project_name=project_name,
        url_override=url_override,
    )

    assert message == messages.CreateAttachmentMessage(
        file_path=original_file,
        file_name=os.path.basename(original_file),
        mime_type="text/plain",
        entity_type="trace",
        entity_id=entity_id,
        project_name=project_name,
        encoded_url_override="aHR0cHM6Ly9leGFtcGxlLmNvbQ==",
    )


def test_attachment_to_message__file_name(original_file: str):
    url_override = "https://example.com"
    entity_id = "123"
    project_name = "test-project"
    attachment_data = attachment.Attachment(
        data=original_file, file_name="test.jpg", create_temp_copy=False
    )

    message = converters.attachment_to_message(
        attachment_data=attachment_data,
        entity_type="trace",
        entity_id=entity_id,
        project_name=project_name,
        url_override=url_override,
    )

    assert message == messages.CreateAttachmentMessage(
        file_path=original_file,
        file_name="test.jpg",
        mime_type="image/jpeg",
        entity_type="trace",
        entity_id=entity_id,
        project_name=project_name,
        encoded_url_override="aHR0cHM6Ly9leGFtcGxlLmNvbQ==",
    )


def test_attachment_to_message__content_type(original_file: str):
    url_override = "https://example.com"
    entity_id = "123"
    project_name = "test-project"
    attachment_data = attachment.Attachment(
        data=original_file, content_type="image/jpeg", create_temp_copy=False
    )

    message = converters.attachment_to_message(
        attachment_data=attachment_data,
        entity_type="trace",
        entity_id=entity_id,
        project_name=project_name,
        url_override=url_override,
    )

    assert message == messages.CreateAttachmentMessage(
        file_path=original_file,
        file_name=os.path.basename(original_file),
        mime_type="image/jpeg",
        entity_type="trace",
        entity_id=entity_id,
        project_name=project_name,
        encoded_url_override="aHR0cHM6Ly9leGFtcGxlLmNvbQ==",
    )


def test_attachment_to_message__create_temp_copy(original_file: str):
    """Test that create_temp_copy creates a temporary file and sets delete_after_upload."""
    url_override = "https://example.com"
    entity_id = "123"
    project_name = "test-project"

    attachment_data = attachment.Attachment(
        data=original_file,
        file_name="test.txt",
        content_type="text/plain",
        create_temp_copy=True,
    )

    message = converters.attachment_to_message(
        attachment_data=attachment_data,
        entity_type="trace",
        entity_id=entity_id,
        project_name=project_name,
        url_override=url_override,
    )

    # The file_path should be different from the original (it's a temp copy)
    assert message.file_path != original_file
    assert message.file_path.endswith(".txt")
    assert os.path.exists(message.file_path)

    # delete_after_upload should be True when create_temp_copy is used
    assert message.delete_after_upload is True

    # Other fields should be preserved
    assert message.file_name == "test.txt"
    assert message.mime_type == "text/plain"
    assert message.entity_type == "trace"
    assert message.entity_id == entity_id
    assert message.project_name == project_name

    # Verify the temp file has the same content
    with open(message.file_path, "r") as f:
        assert f.read() == "test content"

    # Clean up the temp copy
    os.unlink(message.file_path)


def test_attachment_to_message__create_temp_copy_fails_on_open__uses_original_file_and_does_not_delete(
    original_file: str,
):
    """
    Test that when create_temp_copy is True but opening the source file fails,
    the original file is used and delete_after_upload is False.
    """
    url_override = "https://example.com"
    entity_id = "123"
    project_name = "test-project"

    attachment_data = attachment.Attachment(
        data=original_file,
        file_name="test.txt",
        content_type="text/plain",
        create_temp_copy=True,
    )

    # Patch open to fail when trying to read the source file for copying
    original_open = open

    def mock_open_fail_on_rb(file, mode="r", *args, **kwargs):
        if mode == "rb" and file == original_file:
            raise PermissionError("Cannot read source file")
        return original_open(file, mode, *args, **kwargs)

    with mock.patch("builtins.open", side_effect=mock_open_fail_on_rb):
        message = converters.attachment_to_message(
            attachment_data=attachment_data,
            entity_type="trace",
            entity_id=entity_id,
            project_name=project_name,
            url_override=url_override,
        )

    # The file_path should be the original file (fallback behavior)
    assert message.file_path == original_file

    # delete_after_upload should be False to prevent deleting the original file
    assert message.delete_after_upload is False

    # Other fields should be preserved
    assert message.file_name == "test.txt"
    assert message.mime_type == "text/plain"
    assert message.entity_type == "trace"
    assert message.entity_id == entity_id
    assert message.project_name == project_name

    # Original file should still exist
    assert os.path.exists(original_file)

    # Clean up the temp copy
    os.unlink(message.file_path)


def test_attachment_to_message__create_temp_copy_fails_with_delete_after_upload_true__still_does_not_delete(
    original_file: str,
):
    """
    Test that when create_temp_copy fails and delete_after_upload was explicitly True,
    delete_after_upload is still set to False to protect the original file.
    """
    url_override = "https://example.com"
    entity_id = "123"
    project_name = "test-project"

    attachment_data = attachment.Attachment(
        data=original_file,
        file_name="test.txt",
        content_type="text/plain",
        create_temp_copy=True,
    )

    with mock.patch("shutil.copyfileobj", side_effect=IOError("Disk full")):
        message = converters.attachment_to_message(
            attachment_data=attachment_data,
            entity_type="trace",
            entity_id=entity_id,
            project_name=project_name,
            url_override=url_override,
            delete_after_upload=True,  # Explicitly set to True
        )

    # Even though delete_after_upload was True, it should be False after failure
    assert message.delete_after_upload is False
    assert message.file_path == original_file


def test_attachment_to_message__bytes_data():
    """Test that bytes data is written to a temp file and marked for deletion."""
    url_override = "https://example.com"
    entity_id = "123"
    project_name = "test-project"
    data = b"binary content here"

    attachment_data = attachment.Attachment(
        data=data,
        file_name="test.bin",
        content_type="application/octet-stream",
    )

    message = converters.attachment_to_message(
        attachment_data=attachment_data,
        entity_type="trace",
        entity_id=entity_id,
        project_name=project_name,
        url_override=url_override,
    )

    # A temp file should be created
    assert os.path.exists(message.file_path)
    assert message.file_path != "test.bin"

    # delete_after_upload should be True for bytes data
    assert message.delete_after_upload is True

    # Verify file content matches the bytes
    with open(message.file_path, "rb") as f:
        assert f.read() == data

    # Other fields should be preserved
    assert message.file_name == "test.bin"
    assert message.mime_type == "application/octet-stream"
    assert message.entity_type == "trace"
    assert message.entity_id == entity_id
    assert message.project_name == project_name
    assert message.encoded_url_override == "aHR0cHM6Ly9leGFtcGxlLmNvbQ=="

    # Clean up
    os.unlink(message.file_path)


def test_attachment_to_message__bytes_data_without_file_name():
    """Test bytes data without file_name uses temp file basename."""
    url_override = "https://example.com"
    entity_id = "123"
    project_name = "test-project"
    data = b"some binary data"

    attachment_data = attachment.Attachment(
        data=data,
        content_type="image/png",
    )

    message = converters.attachment_to_message(
        attachment_data=attachment_data,
        entity_type="span",
        entity_id=entity_id,
        project_name=project_name,
        url_override=url_override,
    )

    # file_name should be the basename of the temp file
    assert message.file_name == os.path.basename(message.file_path)
    assert message.delete_after_upload is True
    assert message.mime_type == "image/png"
    assert message.entity_type == "span"

    # Clean up
    os.unlink(message.file_path)


def test_attachment_to_message__bytes_data_infers_mime_type_from_file_name():
    """Test that a mime type is inferred from file_name when data is bytes."""
    url_override = "https://example.com"
    entity_id = "123"
    project_name = "test-project"
    data = b"\x00\x01\x02\x03"

    attachment_data = attachment.Attachment(
        data=data,
        file_name="image.png",
    )

    message = converters.attachment_to_message(
        attachment_data=attachment_data,
        entity_type="trace",
        entity_id=entity_id,
        project_name=project_name,
        url_override=url_override,
    )

    assert message.file_name == "image.png"
    assert message.mime_type == "image/png"
    assert message.delete_after_upload is True

    # Clean up
    os.unlink(message.file_path)


def test_attachment_to_message__bytes_data_default_mime_type():
    """Test that bytes data without file_name or content_type use a binary mime type."""
    url_override = "https://example.com"
    entity_id = "123"
    project_name = "test-project"
    data = b"\x00\x01\x02\x03"

    attachment_data = attachment.Attachment(
        data=data,
    )

    message = converters.attachment_to_message(
        attachment_data=attachment_data,
        entity_type="trace",
        entity_id=entity_id,
        project_name=project_name,
        url_override=url_override,
    )

    # Should use default binary mime type
    assert message.mime_type == "application/octet-stream"
    assert message.delete_after_upload is True

    # Clean up
    os.unlink(message.file_path)


def test_attachment_to_message__bytes_data_write_fails__error_reraised():
    """Test behavior when _write_file_like_to_temp_file fails and error reraised."""
    url_override = "https://example.com"
    entity_id = "123"
    project_name = "test-project"
    data = b"test data"

    attachment_data = attachment.Attachment(
        data=data,
        file_name="test.bin",
        content_type="application/octet-stream",
    )

    with mock.patch(
        "opik.api_objects.attachment.converters._write_file_like_to_temp_file",
        side_effect=OSError("Disk full"),
    ):
        with pytest.raises(OSError):
            converters.attachment_to_message(
                attachment_data=attachment_data,
                entity_type="trace",
                entity_id=entity_id,
                project_name=project_name,
                url_override=url_override,
            )


def test_attachment_to_message__base64_string(request):
    """Test that base64-encoded string is decoded and written to temp file."""
    import base64

    url_override = "https://example.com"
    entity_id = "123"
    project_name = "test-project"
    original_content = b"This is test content for base64 encoding"
    base64_string = base64.b64encode(original_content).decode("utf-8")

    attachment_data = attachment.Attachment(
        data=base64_string,
        file_name="test.txt",
        content_type="text/plain",
    )

    message = converters.attachment_to_message(
        attachment_data=attachment_data,
        entity_type="trace",
        entity_id=entity_id,
        project_name=project_name,
        url_override=url_override,
    )

    request.addfinalizer(
        lambda: os.path.exists(message.file_path) and os.unlink(message.file_path)
    )

    # A temp file should be created
    assert os.path.exists(message.file_path)
    assert message.file_path != base64_string

    # delete_after_upload should be True for base64 data
    assert message.delete_after_upload is True

    # Verify file content matches the decoded bytes
    with open(message.file_path, "rb") as f:
        assert f.read() == original_content

    # Other fields should be preserved
    assert message.file_name == "test.txt"
    assert message.mime_type == "text/plain"
    assert message.entity_type == "trace"
    assert message.entity_id == entity_id
    assert message.project_name == project_name


def test_attachment_to_message__invalid_data_raises_error():
    """Test that invalid attachment data (not bytes, file, or base64) raises ValueError."""
    url_override = "https://example.com"
    entity_id = "123"
    project_name = "test-project"

    invalid_data = "not-a-file-and-not-base64!@#$%"

    attachment_data = attachment.Attachment(
        data=invalid_data,
        file_name="test.txt",
    )

    with pytest.raises(
        ValueError,
        match="Attachment data must be bytes, an existing file path, or a valid base64-encoded string",
    ):
        converters.attachment_to_message(
            attachment_data=attachment_data,
            entity_type="trace",
            entity_id=entity_id,
            project_name=project_name,
            url_override=url_override,
        )