blender-tools/blender/tests/smoke_test.py

from __future__ import annotations

import json
import struct
import sys
from pathlib import Path

import bpy


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

import material_id_baker  # noqa: E402


def make_test_object() -> bpy.types.Object:
    mesh = bpy.data.meshes.new("MaterialIDSmokeMesh")
    mesh.from_pydata(
        [
            (-1.0, -1.0, 0.0),
            (0.0, -1.0, 0.0),
            (0.0, 1.0, 0.0),
            (-1.0, 1.0, 0.0),
            (0.0, -1.0, 0.0),
            (1.0, -1.0, 0.0),
            (1.0, 1.0, 0.0),
            (0.0, 1.0, 0.0),
        ],
        [],
        [(0, 1, 2, 3), (4, 5, 6, 7)],
    )
    mesh.update()

    left = bpy.data.materials.new("Left Material")
    right = bpy.data.materials.new("Right Material")
    mesh.materials.append(left)
    mesh.materials.append(right)
    mesh.polygons[0].material_index = 0
    mesh.polygons[1].material_index = 1

    uv_layer = mesh.uv_layers.new(name="ID UV")
    uv_coordinates = (
        (0.05, 0.05),
        (0.45, 0.05),
        (0.45, 0.95),
        (0.05, 0.95),
        (0.55, 0.05),
        (0.95, 0.05),
        (0.95, 0.95),
        (0.55, 0.95),
    )
    for loop, uv in zip(mesh.loops, uv_coordinates, strict=True):
        uv_layer.uv[loop.index].vector = uv

    obj = bpy.data.objects.new("MaterialIDSmokeObject", mesh)
    bpy.context.scene.collection.objects.link(obj)
    obj.modifiers.new(name="Triangulate for evaluated-mesh test", type="TRIANGULATE")
    obj.select_set(True)
    bpy.context.view_layer.objects.active = obj
    return obj


def pixel(image: bpy.types.Image, x: int, y: int) -> tuple[float, float, float, float]:
    offset = (y * image.size[0] + x) * 4
    return tuple(image.pixels[offset : offset + 4])


def close_rgb(actual: tuple[float, ...], expected: list[float], tolerance: float = 0.03) -> bool:
    return all(abs(actual[index] - expected[index]) <= tolerance for index in range(3))


def main() -> None:
    material_id_baker.register()
    source = make_test_object()
    original_engine = bpy.context.scene.render.engine
    bpy.context.scene.cycles.samples = 37
    original_cycles_samples = bpy.context.scene.cycles.samples

    settings = bpy.context.scene.material_id_baker
    assert settings.margin == 1
    settings.resolution = "256"
    settings.margin = 4
    settings.palette_mode = "DISTINCT"
    settings.apply_modifiers = True
    settings.create_export_copy = True
    settings.quick_export_glb = True
    settings.glb_filepath = "/tmp/material_id_baker_smoke_export"
    settings.image_name = "Material_ID_Smoke"
    settings.save_to_disk = True
    settings.filepath = "/tmp/material_id_baker_smoke.png"
    settings.write_legend = True

    result = bpy.ops.object.bake_material_id()
    assert result == {"FINISHED"}, result
    assert bpy.context.scene.render.engine == original_engine
    assert bpy.context.scene.cycles.samples == original_cycles_samples
    assert not any(item.name.startswith("__MID_BAKER_") for item in bpy.data.objects)
    assert not any(item.name.startswith("__MID_BAKER_") for item in bpy.data.materials)

    export_copy = bpy.context.active_object
    assert export_copy is not None
    assert export_copy.name == "MaterialIDSmokeObject_Baked_IDs"
    assert export_copy != source
    assert export_copy.select_get()
    assert not source.select_get()
    assert len(export_copy.data.materials) == 1
    baked_material = export_copy.data.materials[0]
    assert baked_material.name == "baked ids"
    assert all(polygon.material_index == 0 for polygon in export_copy.data.polygons)
    assert len(export_copy.data.polygons) == 4
    assert len(export_copy.modifiers) == 0

    image = bpy.data.images[settings.last_image_name]
    image_nodes = [
        node
        for node in baked_material.node_tree.nodes
        if node.bl_idname == "ShaderNodeTexImage"
    ]
    assert len(image_nodes) == 1
    assert image_nodes[0].image == image
    palette = json.loads(image["material_id_palette"])
    assert len(palette) == 2
    assert palette[0]["material"] == "Left Material"
    assert palette[1]["material"] == "Right Material"
    left_pixel = pixel(image, 64, 128)
    right_pixel = pixel(image, 192, 128)
    assert close_rgb(left_pixel, palette[0]["rgb"]), (left_pixel, palette[0])
    assert close_rgb(right_pixel, palette[1]["rgb"]), (right_pixel, palette[1])
    assert Path("/tmp/material_id_baker_smoke.png").is_file()
    assert Path("/tmp/material_id_baker_smoke.json").is_file()

    glb_path = Path("/tmp/material_id_baker_smoke_export.glb")
    assert glb_path.is_file()
    glb_data = glb_path.read_bytes()
    magic, version, total_length = struct.unpack_from("<4sII", glb_data, 0)
    assert magic == b"glTF"
    assert version == 2
    assert total_length == len(glb_data)
    json_length, json_type = struct.unpack_from("<I4s", glb_data, 12)
    assert json_type == b"JSON"
    glb_json = json.loads(glb_data[20 : 20 + json_length].decode("utf-8"))
    assert len(glb_json["materials"]) == 1
    assert glb_json["materials"][0]["name"] == "baked ids"
    assert len(glb_json["meshes"]) == 1
    assert len(glb_json["nodes"]) == 1
    assert len(glb_json["images"]) == 1
    assert "animations" not in glb_json

    print("MATERIAL_ID_BAKER_SMOKE_TEST_OK")


if __name__ == "__main__":
    main()