79 lines
3.3 KiB
Python
79 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Focused hostile-input tests for the safe motion GGUF converter."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import json
|
|
import struct
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
SPEC = importlib.util.spec_from_file_location(
|
|
"convert_motion_to_gguf", ROOT / "scripts" / "convert_motion_to_gguf.py"
|
|
)
|
|
assert SPEC and SPEC.loader
|
|
CONVERTER = importlib.util.module_from_spec(SPEC)
|
|
sys.modules[SPEC.name] = CONVERTER
|
|
SPEC.loader.exec_module(CONVERTER)
|
|
|
|
|
|
class ConverterInputTest(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.path = Path(self.temp.name) / "input"
|
|
|
|
def tearDown(self) -> None:
|
|
self.temp.cleanup()
|
|
|
|
def write_safe(self, header: str | dict, payload: bytes = b"\0" * 8) -> Path:
|
|
encoded = header.encode() if isinstance(header, str) else json.dumps(header).encode()
|
|
output = self.path.with_suffix(".safetensors")
|
|
output.write_bytes(struct.pack("<Q", len(encoded)) + encoded + payload)
|
|
return output
|
|
|
|
def test_valid_f32_safetensors(self) -> None:
|
|
path = self.write_safe({"weight": {"dtype": "F32", "shape": [2], "data_offsets": [0, 8]}})
|
|
tensors = CONVERTER.read_safetensors(path)
|
|
self.assertEqual([(item.name, item.shape, item.size) for item in tensors], [("weight", (2,), 8)])
|
|
|
|
def test_safetensors_duplicate_json_key_is_rejected(self) -> None:
|
|
path = self.write_safe('{"weight":{"dtype":"F32","shape":[1],"data_offsets":[0,4]},"weight":{"dtype":"F32","shape":[1],"data_offsets":[4,8]}}')
|
|
with self.assertRaises(ValueError):
|
|
CONVERTER.read_safetensors(path)
|
|
|
|
def test_safetensors_out_of_range_and_overlap_are_rejected(self) -> None:
|
|
out_of_range = self.write_safe({"weight": {"dtype": "F32", "shape": [2], "data_offsets": [0, 8]}}, b"\0" * 4)
|
|
with self.assertRaises(ValueError):
|
|
CONVERTER.read_safetensors(out_of_range)
|
|
overlap = self.write_safe({
|
|
"left": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]},
|
|
"right": {"dtype": "F32", "shape": [1], "data_offsets": [2, 6]},
|
|
})
|
|
with self.assertRaises(ValueError):
|
|
CONVERTER.read_safetensors(overlap)
|
|
|
|
def write_npy(self, descriptor: str, payload: bytes) -> Path:
|
|
output = self.path.with_suffix(".npy")
|
|
header = descriptor.encode("latin1")
|
|
output.write_bytes(b"\x93NUMPY" + bytes((1, 0)) + struct.pack("<H", len(header)) + header + payload)
|
|
return output
|
|
|
|
def test_npy_payload_and_layout_are_checked(self) -> None:
|
|
valid = self.write_npy("{'descr': '<f4', 'fortran_order': False, 'shape': (2,), }", b"\0" * 8)
|
|
self.assertEqual(CONVERTER.read_npy(valid, "stats.test").size, 8)
|
|
truncated = self.write_npy("{'descr': '<f4', 'fortran_order': False, 'shape': (2,), }", b"\0" * 4)
|
|
with self.assertRaises(ValueError):
|
|
CONVERTER.read_npy(truncated, "stats.test")
|
|
fortran = self.write_npy("{'descr': '<f4', 'fortran_order': True, 'shape': (1,), }", b"\0" * 4)
|
|
with self.assertRaises(ValueError):
|
|
CONVERTER.read_npy(fortran, "stats.test")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|