import unittest
import tempfile
from pathlib import Path
import torch
from safetensors.torch import save_file

from tools.maya_export.common import iter_ckpt_tensors


class TestCommon(unittest.TestCase):
    def test_iter_ckpt_tensors_skips_prefix(self):
        # tiny safetensors file in a temp dir with keys backbone.lm.model.layers.0.w and depth.x;
        # skip backbone.lm.model.layers. -> only depth.x yielded, once
        with tempfile.TemporaryDirectory() as td:
            p = Path(td) / "model.safetensors"
            tensors = {
                "backbone.lm.model.layers.0.w": torch.tensor([1.0, 2.0]),
                "depth.x": torch.tensor([3.0, 4.0, 5.0]),
                "model.depth.y": torch.tensor([6.0]),
            }
            save_file(tensors, str(p))

            res = list(iter_ckpt_tensors(skip_prefixes=("backbone.lm.model.layers.",), files=[str(p)]))
            names = [name for name, _ in res]
            self.assertIn("depth.x", names)
            self.assertIn("depth.y", names)  # model. prefix stripped
            self.assertNotIn("backbone.lm.model.layers.0.w", names)
            self.assertEqual(len(res), 2)


if __name__ == "__main__":
    unittest.main()