tests/test_preprocessor.py

import tempfile
from pathlib import Path

import numpy as np
import pytest

from server.preprocessor import (
    normalize_shape,
    sample_surface_points,
    voxelize,
    preprocess_from_code,
    _occ_to_trimesh,
)


class TestNormalizeShape:
    def test_centered_at_origin(self, flat_plate_shape):
        normalized, info = normalize_shape(flat_plate_shape)
        bb = normalized.BoundingBox()
        cx = (bb.xmin + bb.xmax) / 2
        cy = (bb.ymin + bb.ymax) / 2
        cz = (bb.zmin + bb.zmax) / 2
        assert abs(cx) < 0.02
        assert abs(cy) < 0.02
        assert abs(cz) < 0.02

    def test_longest_axis_is_x(self, flat_plate_shape):
        normalized, info = normalize_shape(flat_plate_shape)
        bb = normalized.BoundingBox()
        assert bb.xlen >= bb.ylen
        assert bb.xlen >= bb.zlen

    def test_preserves_volume(self, flat_plate_shape):
        original_vol = flat_plate_shape.Volume()
        normalized, _ = normalize_shape(flat_plate_shape)
        new_vol = normalized.Volume()
        assert abs(original_vol - new_vol) / original_vol < 0.01

    def test_transform_info(self, flat_plate_shape):
        _, info = normalize_shape(flat_plate_shape)
        assert "units" in info
        assert info["units"] == "mm"
        assert "translation_to_origin" in info
        assert "longest_axis" in info
        assert info["longest_axis"] == "X"

    def test_box_with_hole_normalized(self, box_with_hole_shape):
        normalized, info = normalize_shape(box_with_hole_shape)
        bb = normalized.BoundingBox()
        assert bb.xlen >= bb.ylen >= bb.zlen


class TestOccToTrimesh:
    def test_returns_trimesh(self, flat_plate_shape):
        normalized, _ = normalize_shape(flat_plate_shape)
        mesh = _occ_to_trimesh(normalized)
        import trimesh
        assert isinstance(mesh, trimesh.Trimesh)
        assert len(mesh.vertices) > 0
        assert len(mesh.faces) > 0

    def test_mesh_is_valid(self, flat_plate_shape):
        normalized, _ = normalize_shape(flat_plate_shape)
        mesh = _occ_to_trimesh(normalized)
        assert not mesh.is_empty


class TestSampleSurfacePoints:
    def test_returns_correct_shape(self, flat_plate_shape):
        normalized, _ = normalize_shape(flat_plate_shape)
        points = sample_surface_points(normalized, n_points=1024)
        assert points.shape == (1024, 3)
        assert points.dtype == np.float32

    def test_points_near_surface(self, flat_plate_shape):
        normalized, _ = normalize_shape(flat_plate_shape)
        bb = normalized.BoundingBox()
        points = sample_surface_points(normalized, n_points=2048)
        assert np.all(points[:, 0] >= bb.xmin - 0.1)
        assert np.all(points[:, 0] <= bb.xmax + 0.1)
        assert np.all(points[:, 1] >= bb.ymin - 0.1)
        assert np.all(points[:, 1] <= bb.ymax + 0.1)
        assert np.all(points[:, 2] >= bb.zmin - 0.1)
        assert np.all(points[:, 2] <= bb.zmax + 0.1)

    def test_different_n_points(self, flat_plate_shape):
        normalized, _ = normalize_shape(flat_plate_shape)
        for n in [256, 512, 2048]:
            pts = sample_surface_points(normalized, n_points=n)
            assert pts.shape[0] == n


class TestVoxelize:
    def test_returns_correct_shape(self, flat_plate_shape):
        normalized, _ = normalize_shape(flat_plate_shape)
        grid = voxelize(normalized, resolution=32)
        assert grid.shape == (32, 32, 32)
        assert grid.dtype == bool

    def test_has_filled_voxels(self, flat_plate_shape):
        normalized, _ = normalize_shape(flat_plate_shape)
        grid = voxelize(normalized, resolution=32)
        assert grid.sum() > 0

    def test_not_all_filled(self, flat_plate_shape):
        normalized, _ = normalize_shape(flat_plate_shape)
        grid = voxelize(normalized, resolution=32)
        assert grid.sum() < 32 * 32 * 32

    def test_hollow_shape_has_fewer_voxels(self, flat_plate_shape, box_with_hole_shape):
        norm_plate, _ = normalize_shape(flat_plate_shape)
        norm_hole, _ = normalize_shape(box_with_hole_shape)
        grid_plate = voxelize(norm_plate, resolution=32)
        grid_hole = voxelize(norm_hole, resolution=32)
        plate_fill = grid_plate.sum() / (32**3)
        hole_fill = grid_hole.sum() / (32**3)
        assert plate_fill > 0
        assert hole_fill > 0

    def test_resolution_64(self, flat_plate_shape):
        normalized, _ = normalize_shape(flat_plate_shape)
        grid = voxelize(normalized, resolution=64)
        assert grid.shape == (64, 64, 64)
        assert grid.sum() > 0


class TestPreprocessFromCode:
    def test_generates_all_files(self, flat_plate_code):
        with tempfile.TemporaryDirectory() as tmpdir:
            gt = preprocess_from_code(flat_plate_code, tmpdir, task_id="test_flat_plate")
            output_path = Path(tmpdir)
            assert (output_path / "ground_truth.json").exists()
            assert (output_path / "ground_truth.step").exists()
            assert (output_path / "ground_truth_normalized.step").exists()
            assert (output_path / "surface_points.npy").exists()
            assert (output_path / "voxels_64.npy").exists()

    def test_ground_truth_json_contents(self, flat_plate_code):
        with tempfile.TemporaryDirectory() as tmpdir:
            gt = preprocess_from_code(flat_plate_code, tmpdir, task_id="test_flat_plate")
            assert "volume_mm3" in gt
            assert "bbox_mm" in gt
            assert "face_count" in gt
            assert "dominant_face_type" in gt
            assert "surface_points_file" in gt
            assert "voxels_file" in gt
            assert gt["volume_mm3"] > 0
            assert gt["face_count"] == 6

    def test_surface_points_file(self, flat_plate_code):
        with tempfile.TemporaryDirectory() as tmpdir:
            preprocess_from_code(flat_plate_code, tmpdir)
            pts = np.load(str(Path(tmpdir) / "surface_points.npy"))
            assert pts.shape == (2048, 3)

    def test_voxels_file(self, flat_plate_code):
        with tempfile.TemporaryDirectory() as tmpdir:
            preprocess_from_code(flat_plate_code, tmpdir)
            vox = np.load(str(Path(tmpdir) / "voxels_64.npy"))
            assert vox.shape == (64, 64, 64)
            assert vox.sum() > 0

    def test_invalid_code_raises(self, invalid_code):
        with tempfile.TemporaryDirectory() as tmpdir:
            with pytest.raises(Exception):
                preprocess_from_code(invalid_code, tmpdir)

    def test_no_result_raises(self, no_result_code):
        with tempfile.TemporaryDirectory() as tmpdir:
            with pytest.raises(ValueError, match="result"):
                preprocess_from_code(no_result_code, tmpdir)