1  import pytest
  2  
  3  from mlflow.tracing.attachments import Attachment
  4  
  5  
  6  def test_attachment_init():
  7      att = Attachment(content_type="image/png", content_bytes=b"fakepng")
  8      assert att.id is not None
  9      assert att.content_type == "image/png"
 10      assert att.content_bytes == b"fakepng"
 11  
 12  
 13  def test_attachment_ids_are_unique():
 14      a = Attachment(content_type="image/png", content_bytes=b"a")
 15      b = Attachment(content_type="image/png", content_bytes=b"b")
 16      assert a.id != b.id
 17  
 18  
 19  def test_attachment_from_file(tmp_path):
 20      file = tmp_path / "test.png"
 21      file.write_bytes(b"fakepng")
 22      att = Attachment.from_file(file)
 23      assert att.content_type == "image/png"
 24      assert att.content_bytes == b"fakepng"
 25  
 26  
 27  def test_attachment_from_file_explicit_content_type(tmp_path):
 28      file = tmp_path / "test.bin"
 29      file.write_bytes(b"data")
 30      att = Attachment.from_file(file, content_type="custom/type")
 31      assert att.content_type == "custom/type"
 32  
 33  
 34  @pytest.mark.parametrize(
 35      ("suffix", "expected"),
 36      [
 37          (".png", "image/png"),
 38          (".jpg", "image/jpeg"),
 39          (".jpeg", "image/jpeg"),
 40          (".mp3", "audio/mpeg"),
 41          (".pdf", "application/pdf"),
 42          (".xyz_unknown", "application/octet-stream"),
 43      ],
 44  )
 45  def test_attachment_from_file_mime_inference(suffix, expected, tmp_path):
 46      file = tmp_path / f"test{suffix}"
 47      file.write_bytes(b"data")
 48      att = Attachment.from_file(file)
 49      assert att.content_type == expected
 50  
 51  
 52  def test_attachment_ref():
 53      att = Attachment(content_type="image/png", content_bytes=b"data")
 54      ref = att.ref("tr-abc123")
 55      assert ref.startswith("mlflow-attachment://")
 56      assert att.id in ref
 57      parsed = Attachment.parse_ref(ref)
 58      assert parsed["content_type"] == "image/png"
 59      assert parsed["trace_id"] == "tr-abc123"
 60      assert parsed["size"] == 4
 61  
 62  
 63  def test_parse_ref_valid():
 64      att = Attachment(content_type="audio/wav", content_bytes=b"data")
 65      ref = att.ref("tr-xyz")
 66      parsed = Attachment.parse_ref(ref)
 67      assert parsed is not None
 68      assert parsed["attachment_id"] == att.id
 69      assert parsed["content_type"] == "audio/wav"
 70      assert parsed["trace_id"] == "tr-xyz"
 71      assert parsed["size"] == 4
 72  
 73  
 74  def test_parse_ref_invalid():
 75      assert Attachment.parse_ref("https://example.com") is None
 76      assert Attachment.parse_ref("not-a-uri") is None
 77      assert Attachment.parse_ref("mlflow-attachment://") is None
 78  
 79  
 80  def test_from_file_nonexistent_path():
 81      with pytest.raises(FileNotFoundError, match="nonexistent"):
 82          Attachment.from_file("/nonexistent/path/image.png")
 83  
 84  
 85  def test_ref_roundtrips_special_characters_in_content_type():
 86      att = Attachment(content_type="application/vnd.custom+json", content_bytes=b"data")
 87      ref = att.ref("tr-123")
 88      parsed = Attachment.parse_ref(ref)
 89      assert parsed["content_type"] == "application/vnd.custom+json"
 90      assert parsed["size"] == 4
 91  
 92  
 93  def test_parse_ref_without_size():
 94      uri = "mlflow-attachment://abc-123?content_type=image%2Fpng&trace_id=tr-456"
 95      parsed = Attachment.parse_ref(uri)
 96      assert parsed is not None
 97      assert parsed["attachment_id"] == "abc-123"
 98      assert parsed["content_type"] == "image/png"
 99      assert parsed["trace_id"] == "tr-456"
100      assert parsed["size"] is None
101  
102  
103  def test_ref_size_reflects_content_length():
104      content = b"x" * 1024
105      att = Attachment(content_type="image/png", content_bytes=content)
106      ref = att.ref("tr-001")
107      parsed = Attachment.parse_ref(ref)
108      assert parsed["size"] == 1024
109  
110  
111  @pytest.mark.parametrize("bad_size", ["abc", "-1", "-100", "0"])
112  def test_parse_ref_invalid_size(bad_size):
113      uri = f"mlflow-attachment://abc-123?content_type=image%2Fpng&trace_id=tr-456&size={bad_size}"
114      parsed = Attachment.parse_ref(uri)
115      assert parsed is not None
116      assert parsed["size"] is None