wehub-resource-sync
117 lines
3.7 KiB
Python
117 lines
3.7 KiB
Python
import pytest
|
|
|
|
from mlflow.tracing.attachments import Attachment
|
|
|
|
|
|
def test_attachment_init():
|
|
att = Attachment(content_type="image/png", content_bytes=b"fakepng")
|
|
assert att.id is not None
|
|
assert att.content_type == "image/png"
|
|
assert att.content_bytes == b"fakepng"
|
|
|
|
|
|
def test_attachment_ids_are_unique():
|
|
a = Attachment(content_type="image/png", content_bytes=b"a")
|
|
b = Attachment(content_type="image/png", content_bytes=b"b")
|
|
assert a.id != b.id
|
|
|
|
|
|
def test_attachment_from_file(tmp_path):
|
|
file = tmp_path / "test.png"
|
|
file.write_bytes(b"fakepng")
|
|
att = Attachment.from_file(file)
|
|
assert att.content_type == "image/png"
|
|
assert att.content_bytes == b"fakepng"
|
|
|
|
|
|
def test_attachment_from_file_explicit_content_type(tmp_path):
|
|
file = tmp_path / "test.bin"
|
|
file.write_bytes(b"data")
|
|
att = Attachment.from_file(file, content_type="custom/type")
|
|
assert att.content_type == "custom/type"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("suffix", "expected"),
|
|
[
|
|
(".png", "image/png"),
|
|
(".jpg", "image/jpeg"),
|
|
(".jpeg", "image/jpeg"),
|
|
(".mp3", "audio/mpeg"),
|
|
(".pdf", "application/pdf"),
|
|
(".xyz_unknown", "application/octet-stream"),
|
|
],
|
|
)
|
|
def test_attachment_from_file_mime_inference(suffix, expected, tmp_path):
|
|
file = tmp_path / f"test{suffix}"
|
|
file.write_bytes(b"data")
|
|
att = Attachment.from_file(file)
|
|
assert att.content_type == expected
|
|
|
|
|
|
def test_attachment_ref():
|
|
att = Attachment(content_type="image/png", content_bytes=b"data")
|
|
ref = att.ref("tr-abc123")
|
|
assert ref.startswith("mlflow-attachment://")
|
|
assert att.id in ref
|
|
parsed = Attachment.parse_ref(ref)
|
|
assert parsed["content_type"] == "image/png"
|
|
assert parsed["trace_id"] == "tr-abc123"
|
|
assert parsed["size"] == 4
|
|
|
|
|
|
def test_parse_ref_valid():
|
|
att = Attachment(content_type="audio/wav", content_bytes=b"data")
|
|
ref = att.ref("tr-xyz")
|
|
parsed = Attachment.parse_ref(ref)
|
|
assert parsed is not None
|
|
assert parsed["attachment_id"] == att.id
|
|
assert parsed["content_type"] == "audio/wav"
|
|
assert parsed["trace_id"] == "tr-xyz"
|
|
assert parsed["size"] == 4
|
|
|
|
|
|
def test_parse_ref_invalid():
|
|
assert Attachment.parse_ref("https://example.com") is None
|
|
assert Attachment.parse_ref("not-a-uri") is None
|
|
assert Attachment.parse_ref("mlflow-attachment://") is None
|
|
|
|
|
|
def test_from_file_nonexistent_path():
|
|
with pytest.raises(FileNotFoundError, match="nonexistent"):
|
|
Attachment.from_file("/nonexistent/path/image.png")
|
|
|
|
|
|
def test_ref_roundtrips_special_characters_in_content_type():
|
|
att = Attachment(content_type="application/vnd.custom+json", content_bytes=b"data")
|
|
ref = att.ref("tr-123")
|
|
parsed = Attachment.parse_ref(ref)
|
|
assert parsed["content_type"] == "application/vnd.custom+json"
|
|
assert parsed["size"] == 4
|
|
|
|
|
|
def test_parse_ref_without_size():
|
|
uri = "mlflow-attachment://abc-123?content_type=image%2Fpng&trace_id=tr-456"
|
|
parsed = Attachment.parse_ref(uri)
|
|
assert parsed is not None
|
|
assert parsed["attachment_id"] == "abc-123"
|
|
assert parsed["content_type"] == "image/png"
|
|
assert parsed["trace_id"] == "tr-456"
|
|
assert parsed["size"] is None
|
|
|
|
|
|
def test_ref_size_reflects_content_length():
|
|
content = b"x" * 1024
|
|
att = Attachment(content_type="image/png", content_bytes=content)
|
|
ref = att.ref("tr-001")
|
|
parsed = Attachment.parse_ref(ref)
|
|
assert parsed["size"] == 1024
|
|
|
|
|
|
@pytest.mark.parametrize("bad_size", ["abc", "-1", "-100", "0"])
|
|
def test_parse_ref_invalid_size(bad_size):
|
|
uri = f"mlflow-attachment://abc-123?content_type=image%2Fpng&trace_id=tr-456&size={bad_size}"
|
|
parsed = Attachment.parse_ref(uri)
|
|
assert parsed is not None
|
|
assert parsed["size"] is None
|