comet-ml--opik/sdks/python/tests/unit/cli/test_attachments.py

"""Tests for attachment support in CLI export and import."""

import json
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock, patch


from opik.cli.exports.project import (
    _download_attachment_file,
    _fetch_attachments,
    export_traces,
)
from opik.cli.imports.project import (
    _upload_attachments_for_trace,
    import_traces_from_directory,
)


# ─────────────────────────────────────────────────────────────────────────────
# Helpers shared across test classes
# ─────────────────────────────────────────────────────────────────────────────


def _rest_attachment(file_name, mime_type="image/png", file_size=1024):
    """Simulate a REST API Attachment response object."""
    return SimpleNamespace(
        file_name=file_name,
        mime_type=mime_type,
        file_size=file_size,
        link="http://storage.example.com/presigned-url",
    )


def _make_attachment_client(trace_attachments=None, span_attachments=None):
    """Build a mock AttachmentClient with configurable attachment lists.

    Args:
        trace_attachments: list of dicts with keys accepted by _rest_attachment()
        span_attachments: mapping of {span_id: [dicts]} for span-level attachments
    """
    client = MagicMock()

    def _get_list(project_name, entity_id, entity_type):
        if entity_type == "trace":
            return [_rest_attachment(**a) for a in (trace_attachments or [])]
        elif entity_type == "span":
            span_atts = (span_attachments or {}).get(entity_id, [])
            return [_rest_attachment(**a) for a in span_atts]
        return []

    client.get_attachment_list.side_effect = _get_list
    client.download_attachment.return_value = iter([b"binary content"])
    return client


def _make_trace(trace_id, name="test-trace"):
    """Simulate a trace object as returned by the REST API."""
    t = MagicMock()
    t.id = trace_id
    t.name = name
    t.model_dump.return_value = {
        "id": trace_id,
        "name": name,
        "start_time": "2026-01-01T00:00:00Z",
        "end_time": "2026-01-01T00:00:01Z",
        "input": {},
        "output": {},
        "metadata": None,
        "tags": None,
        "feedback_scores": None,
        "error_info": None,
        "usage": None,
    }
    return t


def _make_span(span_id, trace_id, parent_span_id=None):
    """Simulate a span object."""
    s = MagicMock()
    s.id = span_id
    s.model_dump.return_value = {
        "id": span_id,
        "trace_id": trace_id,
        "parent_span_id": parent_span_id,
        "name": f"span-{span_id}",
        "type": "general",
        "start_time": "2026-01-01T00:00:00Z",
        "end_time": "2026-01-01T00:00:01Z",
        "input": {},
        "output": {},
        "metadata": None,
        "tags": None,
        "feedback_scores": None,
        "usage": None,
    }
    return s


def _make_opik_client(traces, spans_by_trace_id, attachment_client=None):
    """Build a mock opik.Opik client for export tests.

    Args:
        traces: list of trace mock objects
        spans_by_trace_id: mapping of {trace_id: [span_mock]}
        attachment_client: optional mock AttachmentClient
    """
    client = MagicMock()

    # Paginated traces response — first call returns traces, second returns empty
    page1 = MagicMock()
    page1.content = traces
    page1.total = len(traces)

    page2 = MagicMock()
    page2.content = []
    page2.total = len(traces)

    client.rest_client.traces.get_traces_by_project.side_effect = [page1, page2]

    # Bulk span fetch (Phase 2 in 3-phase export): collect all spans and set trace_id
    all_spans = []
    for tid, spans in spans_by_trace_id.items():
        for span in spans:
            span.trace_id = tid
            all_spans.append(span)

    span_page1 = MagicMock()
    span_page1.content = all_spans
    span_page2 = MagicMock()
    span_page2.content = []

    client.rest_client.spans.get_spans_by_project.side_effect = [span_page1, span_page2]

    client.get_attachment_client.return_value = attachment_client or MagicMock()
    return client


def _write_trace_file(project_dir: Path, trace_data: dict) -> Path:
    """Write a trace JSON file and return its path."""
    project_dir.mkdir(parents=True, exist_ok=True)
    trace_id = trace_data["trace"]["id"]
    path = project_dir / f"trace_{trace_id}.json"
    path.write_text(json.dumps(trace_data))
    return path


def _write_attachment_file(
    project_dir: Path,
    entity_type: str,
    entity_id: str,
    file_name: str,
    content: bytes = b"file data",
) -> Path:
    """Write a binary attachment file and return its path."""
    path = project_dir / "attachments" / entity_type / entity_id / file_name
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(content)
    return path


# ─────────────────────────────────────────────────────────────────────────────
# Tests: _fetch_attachments()
# ─────────────────────────────────────────────────────────────────────────────


class TestFetchAttachments:
    def test_fetch_attachments__trace_entity__returns_metadata(self):
        att_client = _make_attachment_client(
            trace_attachments=[
                {"file_name": "img.png", "mime_type": "image/png", "file_size": 500}
            ]
        )
        result = _fetch_attachments(att_client, "proj", "trace-1", [])

        assert len(result) == 1
        assert result[0] == {
            "entity_type": "trace",
            "entity_id": "trace-1",
            "file_name": "img.png",
            "mime_type": "image/png",
            "file_size": 500,
        }

    def test_fetch_attachments__span_entity__returns_metadata(self):
        att_client = _make_attachment_client(
            span_attachments={
                "span-1": [
                    {"file_name": "data.csv", "mime_type": "text/csv", "file_size": 200}
                ]
            }
        )
        result = _fetch_attachments(att_client, "proj", "trace-1", ["span-1"])

        assert len(result) == 1
        assert result[0] == {
            "entity_type": "span",
            "entity_id": "span-1",
            "file_name": "data.csv",
            "mime_type": "text/csv",
            "file_size": 200,
        }

    def test_fetch_attachments__trace_and_span__returns_both(self):
        att_client = _make_attachment_client(
            trace_attachments=[{"file_name": "trace.png"}],
            span_attachments={
                "span-1": [{"file_name": "span.txt", "mime_type": "text/plain"}]
            },
        )
        result = _fetch_attachments(att_client, "proj", "trace-1", ["span-1"])

        assert len(result) == 2
        types = {a["entity_type"] for a in result}
        assert types == {"trace", "span"}

    def test_fetch_attachments__no_attachments__returns_empty_list(self):
        att_client = _make_attachment_client()
        result = _fetch_attachments(att_client, "proj", "trace-1", ["span-1"])
        assert result == []

    def test_fetch_attachments__multiple_spans__collects_all(self):
        att_client = _make_attachment_client(
            span_attachments={
                "span-a": [{"file_name": "a.png"}],
                "span-b": [{"file_name": "b.png"}, {"file_name": "c.png"}],
            }
        )
        result = _fetch_attachments(att_client, "proj", "trace-1", ["span-a", "span-b"])

        assert len(result) == 3
        file_names = {a["file_name"] for a in result}
        assert file_names == {"a.png", "b.png", "c.png"}

    def test_fetch_attachments__conflict_409__skips_entity(self):
        from opik.rest_api.core.api_error import ApiError

        att_client = MagicMock()
        att_client.get_attachment_list.side_effect = ApiError(
            status_code=409, body="conflict"
        )

        # 409 is silently skipped; the function returns what it could collect
        result = _fetch_attachments(att_client, "proj", "trace-1", ["span-1"])
        assert result == []

    def test_fetch_attachments__unexpected_api_error__raises(self):
        from opik.rest_api.core.api_error import ApiError

        att_client = MagicMock()
        att_client.get_attachment_list.side_effect = ApiError(
            status_code=400, body="bad request"
        )

        import pytest

        with pytest.raises(ApiError):
            _fetch_attachments(att_client, "proj", "trace-1", [])

    def test_fetch_attachments__trace_and_spans__calls_get_list_for_each(self):
        att_client = _make_attachment_client()
        _fetch_attachments(att_client, "proj", "trace-1", ["span-a", "span-b"])

        calls = att_client.get_attachment_list.call_args_list
        # One call for the trace, one per span
        assert len(calls) == 3
        entity_types = [c[0][2] for c in calls]  # 3rd positional arg
        assert entity_types.count("trace") == 1
        assert entity_types.count("span") == 2


# ─────────────────────────────────────────────────────────────────────────────
# Tests: _download_attachment_file()
# ─────────────────────────────────────────────────────────────────────────────


class TestDownloadAttachmentFile:
    _ATT = {
        "entity_type": "trace",
        "entity_id": "trace-123",
        "file_name": "img.png",
        "mime_type": "image/png",
        "file_size": 12,
    }

    def test_download_attachment_file__single_chunk__writes_bytes_to_correct_path(
        self, tmp_path
    ):
        att_client = MagicMock()
        att_client.download_attachment.return_value = iter([b"hello world"])

        result = _download_attachment_file(
            att_client, "proj", self._ATT, tmp_path, force=False
        )

        assert result is True
        dest = tmp_path / "attachments" / "trace" / "trace-123" / "img.png"
        assert dest.exists()
        assert dest.read_bytes() == b"hello world"

    def test_download_attachment_file__multiple_chunks__concatenates_bytes(
        self, tmp_path
    ):
        att_client = MagicMock()
        att_client.download_attachment.return_value = iter(
            [b"chunk1", b"chunk2", b"chunk3"]
        )

        _download_attachment_file(att_client, "proj", self._ATT, tmp_path, force=False)

        dest = tmp_path / "attachments" / "trace" / "trace-123" / "img.png"
        assert dest.read_bytes() == b"chunk1chunk2chunk3"

    def test_download_attachment_file__existing_file_no_force__skips_download(
        self, tmp_path
    ):
        att_client = MagicMock()
        dest = tmp_path / "attachments" / "trace" / "trace-123" / "img.png"
        dest.parent.mkdir(parents=True)
        dest.write_bytes(b"original")

        result = _download_attachment_file(
            att_client, "proj", self._ATT, tmp_path, force=False
        )

        assert result is True
        att_client.download_attachment.assert_not_called()
        assert dest.read_bytes() == b"original"  # unchanged

    def test_download_attachment_file__existing_file_with_force__overwrites(
        self, tmp_path
    ):
        att_client = MagicMock()
        att_client.download_attachment.return_value = iter([b"new content"])
        dest = tmp_path / "attachments" / "trace" / "trace-123" / "img.png"
        dest.parent.mkdir(parents=True)
        dest.write_bytes(b"old content")

        result = _download_attachment_file(
            att_client, "proj", self._ATT, tmp_path, force=True
        )

        assert result is True
        att_client.download_attachment.assert_called_once()
        assert dest.read_bytes() == b"new content"

    def test_download_attachment_file__os_error__returns_false(self, tmp_path):
        att_client = MagicMock()
        att_client.download_attachment.side_effect = OSError("S3 unavailable")

        result = _download_attachment_file(
            att_client, "proj", self._ATT, tmp_path, force=False
        )

        assert result is False

    def test_download_attachment_file__nested_entity_id__creates_parent_directories(
        self, tmp_path
    ):
        att_client = MagicMock()
        att_client.download_attachment.return_value = iter([b"data"])
        att = dict(self._ATT, entity_id="a/b/c", file_name="deep.txt")

        _download_attachment_file(att_client, "proj", att, tmp_path, force=False)

        dest = tmp_path / "attachments" / "trace" / "a/b/c" / "deep.txt"
        assert dest.exists()

    def test_download_attachment_file__span_entity__saved_under_span_directory(
        self, tmp_path
    ):
        att_client = MagicMock()
        att_client.download_attachment.return_value = iter([b"span bytes"])
        att = {
            "entity_type": "span",
            "entity_id": "span-xyz",
            "file_name": "output.txt",
            "mime_type": "text/plain",
            "file_size": 10,
        }

        _download_attachment_file(att_client, "proj", att, tmp_path, force=False)

        dest = tmp_path / "attachments" / "span" / "span-xyz" / "output.txt"
        assert dest.exists()
        assert dest.read_bytes() == b"span bytes"


# ─────────────────────────────────────────────────────────────────────────────
# Tests: _upload_attachments_for_trace()
# ─────────────────────────────────────────────────────────────────────────────


class TestUploadAttachmentsForTrace:
    """Unit tests for the _upload_attachments_for_trace helper.

    This private helper is tested directly (rather than solely via
    import_traces_from_directory) because it contains security-critical path
    validation logic (traversal prevention) and ID-translation rules that are
    difficult to observe through the public API without significant fixture
    overhead.  The integration-level behaviour (attachment uploads wired into
    the full import flow) is covered by TestImportProjectsUploadsAttachments.
    """

    def test_uploads_trace_attachment__with_new_trace_id__happyflow(self, tmp_path):
        client = MagicMock()
        project_dir = tmp_path / "my-project"
        _write_attachment_file(
            project_dir, "trace", "old-trace", "img.png", b"image data"
        )

        attachments = [
            {
                "entity_type": "trace",
                "entity_id": "old-trace",
                "file_name": "img.png",
                "mime_type": "image/png",
            }
        ]
        _upload_attachments_for_trace(
            client=client,
            project_dir=project_dir,
            attachments=attachments,
            new_trace_id="new-trace",
            span_id_map={},
            project_name="proj",
        )

        client.queue_attachment_upload.assert_called_once_with(
            entity_type="trace",
            entity_id="new-trace",
            project_name="proj",
            file_path=str(
                project_dir / "attachments" / "trace" / "old-trace" / "img.png"
            ),
            file_name="img.png",
            mime_type="image/png",
        )

    def test_uploads_span_attachment__with_translated_span_id__happyflow(
        self, tmp_path
    ):
        client = MagicMock()
        project_dir = tmp_path / "my-project"
        _write_attachment_file(project_dir, "span", "old-span", "data.csv", b"csv data")

        attachments = [
            {
                "entity_type": "span",
                "entity_id": "old-span",
                "file_name": "data.csv",
                "mime_type": "text/csv",
            }
        ]
        _upload_attachments_for_trace(
            client=client,
            project_dir=project_dir,
            attachments=attachments,
            new_trace_id="new-trace",
            span_id_map={"old-span": "new-span"},
            project_name="proj",
        )

        client.queue_attachment_upload.assert_called_once_with(
            entity_type="span",
            entity_id="new-span",
            project_name="proj",
            file_path=str(
                project_dir / "attachments" / "span" / "old-span" / "data.csv"
            ),
            file_name="data.csv",
            mime_type="text/csv",
        )

    def test_uploads_attachments__trace_and_span__both_uploaded(self, tmp_path):
        client = MagicMock()
        project_dir = tmp_path / "my-project"
        _write_attachment_file(project_dir, "trace", "old-trace", "img.png")
        _write_attachment_file(project_dir, "span", "old-span", "data.csv")

        attachments = [
            {
                "entity_type": "trace",
                "entity_id": "old-trace",
                "file_name": "img.png",
                "mime_type": "image/png",
            },
            {
                "entity_type": "span",
                "entity_id": "old-span",
                "file_name": "data.csv",
                "mime_type": "text/csv",
            },
        ]
        _upload_attachments_for_trace(
            client=client,
            project_dir=project_dir,
            attachments=attachments,
            new_trace_id="new-trace",
            span_id_map={"old-span": "new-span"},
            project_name="proj",
        )

        assert client.queue_attachment_upload.call_count == 2

    def test_upload_attachment__span_not_in_id_map__skips(self, tmp_path):
        client = MagicMock()
        project_dir = tmp_path / "my-project"
        _write_attachment_file(project_dir, "span", "unknown-span", "file.txt")

        attachments = [
            {
                "entity_type": "span",
                "entity_id": "unknown-span",
                "file_name": "file.txt",
                "mime_type": "text/plain",
            }
        ]
        _upload_attachments_for_trace(
            client=client,
            project_dir=project_dir,
            attachments=attachments,
            new_trace_id="new-trace",
            span_id_map={},  # unknown-span not in map
            project_name="proj",
        )

        client.queue_attachment_upload.assert_not_called()

    def test_upload_attachment__file_missing_on_disk__skips(self, tmp_path):
        client = MagicMock()
        project_dir = tmp_path / "my-project"
        # No file written on disk

        attachments = [
            {
                "entity_type": "trace",
                "entity_id": "old-trace",
                "file_name": "missing.png",
                "mime_type": "image/png",
            }
        ]
        _upload_attachments_for_trace(
            client=client,
            project_dir=project_dir,
            attachments=attachments,
            new_trace_id="new-trace",
            span_id_map={},
            project_name="proj",
        )

        client.queue_attachment_upload.assert_not_called()

    def test_upload_attachment__missing_required_keys__skips(self, tmp_path):
        client = MagicMock()
        project_dir = tmp_path / "my-project"

        attachments = [
            {"entity_type": "trace"},  # missing entity_id and file_name
        ]
        _upload_attachments_for_trace(
            client=client,
            project_dir=project_dir,
            attachments=attachments,
            new_trace_id="new-trace",
            span_id_map={},
            project_name="proj",
        )

        client.queue_attachment_upload.assert_not_called()


# ─────────────────────────────────────────────────────────────────────────────
# Integration tests: export_traces() writes attachments to JSON and disk
# ─────────────────────────────────────────────────────────────────────────────


TRACE_ID = "00000000-0000-7000-0000-000000000001"
SPAN_ID = "00000000-0000-7000-0000-000000000002"


class TestExportTracesWritesAttachmentMetadata:
    """Verify that export_traces() embeds attachment metadata in the JSON file."""

    def _run_export(self, client, project_dir, include_attachments=True, force=False):
        export_traces(
            client=client,
            project_name="test-project",
            project_dir=project_dir,
            max_results=None,
            filter_string=None,
            format="json",
            debug=False,
            force=force,
            show_progress=False,
            include_attachments=include_attachments,
        )

    def test_export_traces__with_attachment__key_present_in_json(self, tmp_path):
        trace = _make_trace(TRACE_ID)
        span = _make_span(SPAN_ID, TRACE_ID)
        att_client = _make_attachment_client(
            trace_attachments=[
                {"file_name": "img.png", "mime_type": "image/png", "file_size": 100}
            ]
        )
        client = _make_opik_client(
            traces=[trace],
            spans_by_trace_id={TRACE_ID: [span]},
            attachment_client=att_client,
        )

        self._run_export(client, tmp_path, include_attachments=True)

        trace_file = tmp_path / f"trace_{TRACE_ID}.json"
        assert trace_file.exists()
        data = json.loads(trace_file.read_text())
        assert "attachments" in data
        assert len(data["attachments"]) == 1
        att = data["attachments"][0]
        assert att["entity_type"] == "trace"
        assert att["entity_id"] == TRACE_ID
        assert att["file_name"] == "img.png"
        assert att["mime_type"] == "image/png"

    def test_export_traces__no_attachments__list_is_empty(self, tmp_path):
        trace = _make_trace(TRACE_ID)
        att_client = _make_attachment_client()  # no attachments
        client = _make_opik_client(
            traces=[trace],
            spans_by_trace_id={TRACE_ID: []},
            attachment_client=att_client,
        )

        self._run_export(client, tmp_path, include_attachments=True)

        data = json.loads((tmp_path / f"trace_{TRACE_ID}.json").read_text())
        assert data["attachments"] == []

    def test_export_traces__no_attachments_flag__omits_metadata(self, tmp_path):
        trace = _make_trace(TRACE_ID)
        att_client = _make_attachment_client(
            trace_attachments=[{"file_name": "img.png"}]
        )
        client = _make_opik_client(
            traces=[trace],
            spans_by_trace_id={TRACE_ID: []},
            attachment_client=att_client,
        )

        self._run_export(client, tmp_path, include_attachments=False)

        data = json.loads((tmp_path / f"trace_{TRACE_ID}.json").read_text())
        # attachments key present but empty (include_attachments=False → att_client is None)
        assert data["attachments"] == []
        # And get_attachment_client was never called
        client.get_attachment_client.assert_not_called()

    def test_export_traces__attachment_download__writes_to_correct_path(self, tmp_path):
        trace = _make_trace(TRACE_ID)
        att_client = _make_attachment_client(
            trace_attachments=[
                {"file_name": "photo.jpg", "mime_type": "image/jpeg", "file_size": 50}
            ]
        )
        att_client.download_attachment.return_value = iter([b"jpeg bytes"])
        client = _make_opik_client(
            traces=[trace],
            spans_by_trace_id={TRACE_ID: []},
            attachment_client=att_client,
        )

        self._run_export(client, tmp_path, include_attachments=True)

        dest = tmp_path / "attachments" / "trace" / TRACE_ID / "photo.jpg"
        assert dest.exists()
        assert dest.read_bytes() == b"jpeg bytes"

    def test_export_traces__span_attachment__binary_downloaded(self, tmp_path):
        trace = _make_trace(TRACE_ID)
        span = _make_span(SPAN_ID, TRACE_ID)
        att_client = _make_attachment_client(
            span_attachments={
                SPAN_ID: [
                    {
                        "file_name": "span_output.txt",
                        "mime_type": "text/plain",
                        "file_size": 10,
                    }
                ]
            }
        )
        att_client.download_attachment.return_value = iter([b"span text"])
        client = _make_opik_client(
            traces=[trace],
            spans_by_trace_id={TRACE_ID: [span]},
            attachment_client=att_client,
        )

        self._run_export(client, tmp_path, include_attachments=True)

        dest = tmp_path / "attachments" / "span" / SPAN_ID / "span_output.txt"
        assert dest.exists()
        assert dest.read_bytes() == b"span text"

    def test_export_traces__existing_attachment_file_no_force__skips_download(
        self, tmp_path
    ):
        trace = _make_trace(TRACE_ID)
        att_client = _make_attachment_client(
            trace_attachments=[{"file_name": "img.png"}]
        )
        # Pre-create the attachment file
        dest = tmp_path / "attachments" / "trace" / TRACE_ID / "img.png"
        dest.parent.mkdir(parents=True)
        dest.write_bytes(b"already here")

        client = _make_opik_client(
            traces=[trace],
            spans_by_trace_id={TRACE_ID: []},
            attachment_client=att_client,
        )

        self._run_export(client, tmp_path, include_attachments=True, force=False)

        # File unchanged; download_attachment was never called
        att_client.download_attachment.assert_not_called()
        assert dest.read_bytes() == b"already here"


# ─────────────────────────────────────────────────────────────────────────────
# Integration tests: import_traces_from_directory() uploads attachments
# ─────────────────────────────────────────────────────────────────────────────

_TRACE_WITH_ATTACHMENTS = {
    "trace": {
        "id": "orig-trace-id",
        "name": "test-trace",
        "start_time": "2026-01-01T00:00:00Z",
        "end_time": "2026-01-01T00:00:01Z",
        "input": {},
        "output": {},
        "metadata": None,
        "tags": None,
        "feedback_scores": None,
        "error_info": None,
        "thread_id": None,
    },
    "spans": [
        {
            "id": "orig-span-id",
            "trace_id": "orig-trace-id",
            "parent_span_id": None,
            "name": "child-span",
            "type": "general",
            "start_time": "2026-01-01T00:00:00Z",
            "end_time": "2026-01-01T00:00:00.500000Z",
            "input": {},
            "output": {},
            "metadata": None,
            "tags": None,
            "feedback_scores": None,
            "usage": None,
        }
    ],
    "attachments": [
        {
            "entity_type": "trace",
            "entity_id": "orig-trace-id",
            "file_name": "trace_img.png",
            "mime_type": "image/png",
            "file_size": 100,
        },
        {
            "entity_type": "span",
            "entity_id": "orig-span-id",
            "file_name": "span_data.csv",
            "mime_type": "text/csv",
            "file_size": 50,
        },
    ],
    "downloaded_at": "2026-01-01T00:01:00",
    "project_name": "test-project",
}


def _make_import_client(new_trace_id="new-trace-id", new_span_id="new-span-id"):
    """Build a mock opik.Opik client for import tests."""
    client = MagicMock()
    mock_trace = MagicMock()
    mock_trace.id = new_trace_id
    client.trace.return_value = mock_trace

    mock_span = MagicMock()
    mock_span.id = new_span_id
    client.span.return_value = mock_span

    # flush() is called once per project after all traces are processed.
    client.flush.return_value = True
    return client


class TestImportProjectsUploadsAttachments:
    """Integration tests verifying attachment uploads during project import."""

    def test_import_projects__trace_attachment__uploaded_with_new_trace_id(
        self, tmp_path
    ):
        client = _make_import_client(new_trace_id="new-trace", new_span_id="new-span")
        project_dir = tmp_path / "test-project"
        _write_trace_file(project_dir, _TRACE_WITH_ATTACHMENTS)
        _write_attachment_file(
            project_dir, "trace", "orig-trace-id", "trace_img.png", b"img bytes"
        )
        _write_attachment_file(
            project_dir, "span", "orig-span-id", "span_data.csv", b"csv bytes"
        )

        import_traces_from_directory(
            client=client,
            project_dir=project_dir,
            project_name="test-project",
            dry_run=False,
            name_pattern=None,
            debug=False,
            include_attachments=True,
        )

        upload_calls = client.queue_attachment_upload.call_args_list

        # Verify trace attachment uses new trace ID
        trace_call = next(
            c for c in upload_calls if c.kwargs.get("entity_type") == "trace"
        )
        assert trace_call.kwargs["entity_id"] == "new-trace"
        assert trace_call.kwargs["file_name"] == "trace_img.png"
        assert trace_call.kwargs["mime_type"] == "image/png"

    def test_import_projects__span_attachment__uploaded_with_new_span_id(
        self, tmp_path
    ):
        client = _make_import_client(new_trace_id="new-trace", new_span_id="new-span")
        project_dir = tmp_path / "test-project"
        _write_trace_file(project_dir, _TRACE_WITH_ATTACHMENTS)
        _write_attachment_file(project_dir, "trace", "orig-trace-id", "trace_img.png")
        _write_attachment_file(project_dir, "span", "orig-span-id", "span_data.csv")

        import_traces_from_directory(
            client=client,
            project_dir=project_dir,
            project_name="test-project",
            dry_run=False,
            name_pattern=None,
            debug=False,
            include_attachments=True,
        )

        upload_calls = client.queue_attachment_upload.call_args_list

        span_call = next(
            c for c in upload_calls if c.kwargs.get("entity_type") == "span"
        )
        assert span_call.kwargs["entity_id"] == "new-span"
        assert span_call.kwargs["file_name"] == "span_data.csv"

    def test_import_projects__multiple_attachments__all_uploaded(self, tmp_path):
        client = _make_import_client()
        project_dir = tmp_path / "test-project"
        _write_trace_file(project_dir, _TRACE_WITH_ATTACHMENTS)
        _write_attachment_file(project_dir, "trace", "orig-trace-id", "trace_img.png")
        _write_attachment_file(project_dir, "span", "orig-span-id", "span_data.csv")

        import_traces_from_directory(
            client=client,
            project_dir=project_dir,
            project_name="test-project",
            dry_run=False,
            name_pattern=None,
            debug=False,
            include_attachments=True,
        )

        assert client.queue_attachment_upload.call_count == 2

    def test_import_projects__no_attachments_flag__skips_all_uploads(self, tmp_path):
        client = _make_import_client()
        project_dir = tmp_path / "test-project"
        _write_trace_file(project_dir, _TRACE_WITH_ATTACHMENTS)
        _write_attachment_file(project_dir, "trace", "orig-trace-id", "trace_img.png")

        import_traces_from_directory(
            client=client,
            project_dir=project_dir,
            project_name="test-project",
            dry_run=False,
            name_pattern=None,
            debug=False,
            include_attachments=False,
        )

        # queue_attachment_upload must not be called at all
        client.queue_attachment_upload.assert_not_called()

    def test_import_projects__attachment_files_missing__trace_still_imported(
        self, tmp_path
    ):
        """Missing attachment files emit a warning but don't abort the trace import."""
        client = _make_import_client()
        project_dir = tmp_path / "test-project"
        _write_trace_file(project_dir, _TRACE_WITH_ATTACHMENTS)
        # Do NOT write attachment files — simulate export with --no-attachments

        import_traces_from_directory(
            client=client,
            project_dir=project_dir,
            project_name="test-project",
            dry_run=False,
            name_pattern=None,
            debug=False,
            include_attachments=True,
        )

        # Trace and span must still be created
        assert client.trace.call_count == 1
        assert client.span.call_count == 1
        # No uploads because files don't exist
        client.queue_attachment_upload.assert_not_called()

    def test_import_projects__dry_run__skips_uploads(self, tmp_path):
        client = _make_import_client()
        project_dir = tmp_path / "test-project"
        _write_trace_file(project_dir, _TRACE_WITH_ATTACHMENTS)
        _write_attachment_file(project_dir, "trace", "orig-trace-id", "trace_img.png")

        import_traces_from_directory(
            client=client,
            project_dir=project_dir,
            project_name="test-project",
            dry_run=True,
            name_pattern=None,
            debug=False,
            include_attachments=True,
        )

        client.queue_attachment_upload.assert_not_called()

    def test_import_projects__no_attachments_key__imports_normally(self, tmp_path):
        """Trace files from before attachment support (no 'attachments' key) import fine."""
        trace_data = {
            "trace": {
                "id": "legacy-trace",
                "name": "legacy",
                "start_time": "2026-01-01T00:00:00Z",
                "end_time": "2026-01-01T00:00:01Z",
                "input": {},
                "output": {},
                "metadata": None,
                "tags": None,
                "feedback_scores": None,
                "error_info": None,
                "thread_id": None,
            },
            "spans": [],
            "downloaded_at": "2026-01-01T00:01:00",
            "project_name": "test-project",
            # NOTE: no "attachments" key
        }
        client = _make_import_client()
        project_dir = tmp_path / "test-project"
        _write_trace_file(project_dir, trace_data)

        import_traces_from_directory(
            client=client,
            project_dir=project_dir,
            project_name="test-project",
            dry_run=False,
            name_pattern=None,
            debug=False,
            include_attachments=True,
        )

        assert client.trace.call_count == 1
        att_client = client.get_attachment_client.return_value
        att_client.upload_attachment.assert_not_called()


# ─────────────────────────────────────────────────────────────────────────────
# CLI flag tests: --no-attachments is accepted and propagates correctly
# ─────────────────────────────────────────────────────────────────────────────


class TestNoAttachmentsCliFlag:
    """Verify --no-attachments is wired through the CLI commands."""

    def test_export_traces__no_attachments_flag__accepted(self):
        from click.testing import CliRunner
        from opik.cli.exports.project import export_traces_command

        runner = CliRunner()
        with patch("opik.cli.exports.project.export_project_traces") as mock_export:
            result = runner.invoke(
                export_traces_command,
                ["--no-attachments"],
                obj={"workspace": "ws", "project_name": "my-project", "api_key": None},
                catch_exceptions=False,
            )

        assert result.exit_code == 0, result.output
        _, kwargs = mock_export.call_args
        assert kwargs.get("include_attachments") is False

    def test_export_traces__default__includes_attachments(self):
        from click.testing import CliRunner
        from opik.cli.exports.project import export_traces_command

        runner = CliRunner()
        with patch("opik.cli.exports.project.export_project_traces") as mock_export:
            result = runner.invoke(
                export_traces_command,
                [],
                obj={"workspace": "ws", "project_name": "my-project", "api_key": None},
                catch_exceptions=False,
            )

        assert result.exit_code == 0, result.output
        _, kwargs = mock_export.call_args
        assert kwargs.get("include_attachments") is True

    def test_import_traces__no_attachments_flag__accepted(self):
        from click.testing import CliRunner
        from opik.cli.imports import import_group

        runner = CliRunner()
        # Patch _import_by_type where it is *defined* so the import command
        # calls the mock rather than the real function.
        with patch("opik.cli.imports._import_by_type") as mock_import:
            result = runner.invoke(
                import_group,
                ["my-workspace", "my-project", "traces", "--no-attachments"],
                catch_exceptions=False,
            )

        assert result.exit_code == 0, result.output
        assert mock_import.called
        _, kwargs = mock_import.call_args
        assert kwargs.get("include_attachments") is False

    def test_import_traces__default__includes_attachments(self):
        from click.testing import CliRunner
        from opik.cli.imports import import_group

        runner = CliRunner()
        with patch("opik.cli.imports._import_by_type") as mock_import:
            result = runner.invoke(
                import_group,
                ["my-workspace", "my-project", "traces"],
                catch_exceptions=False,
            )

        assert result.exit_code == 0, result.output
        assert mock_import.called
        _, kwargs = mock_import.call_args
        assert kwargs.get("include_attachments") is True