blender-tools/tools/tests/test_combine_substance_textures.py

from __future__ import annotations

import json
import sys
import tempfile
import unittest
from pathlib import Path

from PIL import Image


TOOLS_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(TOOLS_ROOT))

import combine_substance_textures as combine  # noqa: E402


class CombineSubstanceTexturesTest(unittest.TestCase):
    def setUp(self) -> None:
        self.temp_dir = tempfile.TemporaryDirectory(prefix="substance-combine-test-")
        self.root = Path(self.temp_dir.name)
        self.input_dir = self.root / "exports"
        self.output_dir = self.root / "combined"
        self.input_dir.mkdir()

        self.id_map = self.root / "Material_ID.png"
        id_image = Image.new("RGB", (4, 2), (255, 0, 0))
        for y in range(2):
            for x in range(2, 4):
                id_image.putpixel((x, y), (0, 255, 0))
        id_image.save(self.id_map)

        self.legend = self.root / "Material_ID.json"
        self.legend.write_text(
            json.dumps(
                {
                    "materials": [
                        {"slot": 1, "material": "Left Metal", "rgb": [1.0, 0.0, 0.0]},
                        {"slot": 2, "material": "Right Paint", "rgb": [0.0, 1.0, 0.0]},
                    ]
                }
            ),
            encoding="utf-8",
        )

    def tearDown(self) -> None:
        self.temp_dir.cleanup()

    def write_texture(self, name: str, color: tuple[int, int, int]) -> None:
        Image.new("RGB", (4, 2), color).save(self.input_dir / name)

    def test_combines_multiple_channels_by_id_color(self) -> None:
        self.write_texture("Robot_Left_Metal_BaseColor.png", (180, 20, 10))
        self.write_texture("Robot_Right_Paint_BaseColor.png", (10, 40, 220))
        self.write_texture("Robot_Left_Metal_Normal.png", (100, 120, 255))
        self.write_texture("Robot_Right_Paint_Normal.png", (140, 125, 250))

        outputs, unmatched = combine.combine_all(
            id_map_path=self.id_map,
            legend_path=self.legend,
            input_dir=self.input_dir,
            output_dir=self.output_dir,
        )

        self.assertEqual(unmatched, [])
        self.assertEqual(
            {path.name for path in outputs},
            {"Robot_BaseColor.png", "Robot_Normal.png"},
        )
        with Image.open(self.output_dir / "Robot_BaseColor.png") as result:
            self.assertEqual(result.getpixel((0, 0)), (180, 20, 10))
            self.assertEqual(result.getpixel((3, 1)), (10, 40, 220))

    def test_aliases_and_missing_materials(self) -> None:
        self.write_texture("asset_left_roughness.png", (32, 32, 32))
        materials = combine.load_legend(
            self.legend,
            {"left metal": "left", "right paint": "right"},
        )
        groups, unmatched = combine.discover_groups(self.input_dir, materials)
        self.assertEqual(unmatched, [])
        self.assertEqual(len(groups), 1)
        self.assertEqual(groups[0].output_name, "asset_roughness.png")
        self.assertEqual(set(groups[0].textures), {"left metal"})

    def test_rejects_mismatched_resolution(self) -> None:
        Image.new("RGB", (8, 8), (1, 2, 3)).save(
            self.input_dir / "Left_Metal_BaseColor.png"
        )
        with self.assertRaisesRegex(combine.CombineError, "ID map"):
            combine.combine_all(
                id_map_path=self.id_map,
                legend_path=self.legend,
                input_dir=self.input_dir,
                output_dir=self.output_dir,
            )

    def test_rejects_overlapping_tolerant_masks(self) -> None:
        self.legend.write_text(
            json.dumps(
                {
                    "materials": [
                        {"slot": 1, "material": "Left", "rgb": [1 / 255, 0.0, 0.0]},
                        {"slot": 2, "material": "Right", "rgb": [2 / 255, 0.0, 0.0]},
                    ]
                }
            ),
            encoding="utf-8",
        )
        Image.new("RGB", (4, 2), (1, 0, 0)).save(self.id_map)
        self.write_texture("Left_BaseColor.png", (10, 10, 10))
        self.write_texture("Right_BaseColor.png", (20, 20, 20))
        with self.assertRaisesRegex(combine.CombineError, "overlaps"):
            combine.combine_all(
                id_map_path=self.id_map,
                legend_path=self.legend,
                input_dir=self.input_dir,
                output_dir=self.output_dir,
                tolerance=1,
            )

    def test_auto_tolerance_handles_float_to_png_rounding(self) -> None:
        self.legend.write_text(
            json.dumps(
                {
                    "materials": [
                        {"slot": 1, "material": "Left", "rgb": [0.1, 0.2, 1.0]},
                    ]
                }
            ),
            encoding="utf-8",
        )
        # Python rounds 0.1 * 255 to 26, while Blender may store it as 25.
        Image.new("RGB", (4, 2), (25, 51, 255)).save(self.id_map)
        self.write_texture("Left_BaseColor.png", (12, 34, 56))
        outputs, _ = combine.combine_all(
            id_map_path=self.id_map,
            legend_path=self.legend,
            input_dir=self.input_dir,
            output_dir=self.output_dir,
        )
        with Image.open(outputs[0]) as result:
            self.assertEqual(result.getpixel((0, 0)), (12, 34, 56))

    def test_recursive_material_directories_are_grouped(self) -> None:
        left_dir = self.input_dir / "Robot_Left_Metal"
        right_dir = self.input_dir / "Robot_Right_Paint"
        left_dir.mkdir()
        right_dir.mkdir()
        Image.new("L", (4, 2), 40).save(left_dir / "Roughness.png")
        Image.new("L", (4, 2), 210).save(right_dir / "Roughness.png")

        outputs, _ = combine.combine_all(
            id_map_path=self.id_map,
            legend_path=self.legend,
            input_dir=self.input_dir,
            output_dir=self.output_dir,
            recursive=True,
        )
        self.assertEqual(outputs[0].relative_to(self.output_dir), Path("Robot/Roughness.png"))
        with Image.open(outputs[0]) as result:
            self.assertEqual(result.getpixel((0, 0)), 40)
            self.assertEqual(result.getpixel((3, 1)), 210)


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