1  from mlflow.entities.span import LiveSpan, Span
  2  from mlflow.tracing.attachments import Attachment
  3  
  4  
  5  def _make_live_span(trace_id="tr-test123"):
  6      from opentelemetry.sdk.trace import TracerProvider
  7  
  8      tracer = TracerProvider().get_tracer("test")
  9      otel_span = tracer.start_span("test_span")
 10      return LiveSpan(otel_span, trace_id=trace_id)
 11  
 12  
 13  def test_set_inputs_extracts_attachment():
 14      span = _make_live_span()
 15      att = Attachment(content_type="image/png", content_bytes=b"img")
 16      span.set_inputs({"image": att})
 17  
 18      assert att.id in span._attachments
 19      assert span._attachments[att.id] is att
 20      inputs = span.inputs
 21      assert isinstance(inputs["image"], str)
 22      assert inputs["image"].startswith("mlflow-attachment://")
 23      parsed = Attachment.parse_ref(inputs["image"])
 24      assert parsed["trace_id"] == "tr-test123"
 25      assert parsed["attachment_id"] == att.id
 26  
 27  
 28  def test_set_outputs_extracts_attachment():
 29      span = _make_live_span()
 30      att = Attachment(content_type="audio/wav", content_bytes=b"audio")
 31      span.set_outputs({"sound": att})
 32  
 33      assert att.id in span._attachments
 34      outputs = span.outputs
 35      assert isinstance(outputs["sound"], str)
 36      parsed = Attachment.parse_ref(outputs["sound"])
 37      assert parsed["content_type"] == "audio/wav"
 38  
 39  
 40  def test_nested_dict_extraction():
 41      span = _make_live_span()
 42      att = Attachment(content_type="image/jpeg", content_bytes=b"jpg")
 43      span.set_inputs({"nested": {"deep": att}})
 44  
 45      assert att.id in span._attachments
 46      inputs = span.inputs
 47      assert isinstance(inputs["nested"]["deep"], str)
 48      assert att.id in inputs["nested"]["deep"]
 49  
 50  
 51  def test_list_extraction():
 52      span = _make_live_span()
 53      att1 = Attachment(content_type="image/png", content_bytes=b"a")
 54      att2 = Attachment(content_type="image/png", content_bytes=b"b")
 55      span.set_inputs({"images": [att1, att2]})
 56  
 57      assert att1.id in span._attachments
 58      assert att2.id in span._attachments
 59      inputs = span.inputs
 60      assert len(inputs["images"]) == 2
 61      assert all(isinstance(v, str) for v in inputs["images"])
 62  
 63  
 64  def test_mixed_values_only_replaces_attachments():
 65      span = _make_live_span()
 66      att = Attachment(content_type="image/png", content_bytes=b"img")
 67      span.set_inputs({"image": att, "text": "hello", "number": 42})
 68  
 69      assert att.id in span._attachments
 70      inputs = span.inputs
 71      assert inputs["text"] == "hello"
 72      assert inputs["number"] == 42
 73      assert isinstance(inputs["image"], str)
 74  
 75  
 76  def test_non_dict_input_without_attachment():
 77      span = _make_live_span()
 78      span.set_inputs("plain string")
 79      assert span._attachments == {}
 80      assert span.inputs == "plain string"
 81  
 82  
 83  def test_tuple_extraction():
 84      span = _make_live_span()
 85      att = Attachment(content_type="image/png", content_bytes=b"img")
 86      span.set_inputs({"images": (att, "text")})
 87  
 88      assert att.id in span._attachments
 89      inputs = span.inputs
 90      assert isinstance(inputs["images"], list)
 91      assert isinstance(inputs["images"][0], str)
 92      assert inputs["images"][0].startswith("mlflow-attachment://")
 93      assert inputs["images"][1] == "text"
 94  
 95  
 96  def test_set_inputs_twice_accumulates_attachments():
 97      span = _make_live_span()
 98      att1 = Attachment(content_type="image/png", content_bytes=b"first")
 99      att2 = Attachment(content_type="image/jpeg", content_bytes=b"second")
100  
101      span.set_inputs({"img": att1})
102      span.set_outputs({"img": att2})
103  
104      assert att1.id in span._attachments
105      assert att2.id in span._attachments
106      assert len(span._attachments) == 2
107  
108  
109  def test_to_immutable_span_propagates_attachments():
110      span = _make_live_span()
111      att = Attachment(content_type="image/png", content_bytes=b"img")
112      span.set_inputs({"image": att})
113  
114      immutable = span.to_immutable_span()
115      assert isinstance(immutable, Span)
116      assert att.id in immutable._attachments
117      assert immutable._attachments[att.id] is att