Cosplay/nhf/test.py

"""
Unit tests for tooling
"""
from dataclasses import dataclass
import math
import unittest
import cadquery as Cq
from nhf.build import Model, target
from nhf.parts.item import Item
import nhf.checks
import nhf.geometry
import nhf.utils

# Color presets for testing purposes
color_parent = Cq.Color(0.7, 0.7, 0.5, 0.5)
color_child = Cq.Color(0.5, 0.7, 0.7, 0.5)

def makeSphere(r: float) -> Cq.Solid:
    """
    Makes a full sphere. The default function makes a hemisphere
    """
    return Cq.Solid.makeSphere(r, angleDegrees1=-90)

@dataclass(frozen=True)
class MassBall(Item):
    """
    A ball with fixed mass
    """
    radius: float = 0.2

    @property
    def name(self) -> str:
        return f"MassBall {self.mass}"
    def generate(self) -> Cq.Solid:
        return makeSphere(self.radius)


class BuildScaffold(Model):

    def __init__(self):
        super().__init__(name="scaffold")

    @target(name="obj1")
    def o1(self):
        return Cq.Solid.makeBox(10, 10, 10)

    def o2(self):
        return Cq.Solid.makeCylinder(10, 20)

class TestBuild(unittest.TestCase):

    def test_build_scaffold(self):
        s = BuildScaffold()
        names = ["obj1"]
        self.assertEqual(s.target_names, names)
        self.assertEqual(s.check_all(), len(names))

class TestChecks(unittest.TestCase):

    def intersect_test_case(self, offset):
        assembly = (
            Cq.Assembly()
            .add(Cq.Solid.makeBox(10, 10, 10),
                 name="c1",
                 loc=Cq.Location((0, 0, 0)))
            .add(Cq.Solid.makeBox(10, 10, 10),
                 name="c2",
                 loc=Cq.Location((0, 0, offset)))
        )
        coll = nhf.checks.pairwise_intersection(assembly)
        if -10 < offset and offset < 10:
            self.assertEqual(len(coll), 1)
        else:
            self.assertEqual(coll, [])

    def test_intersect(self):
        for offset in [9, 10, 11, -10]:
            with self.subTest(offset=offset):
                self.intersect_test_case(offset)

class TestGeometry(unittest.TestCase):

    def test_min_radius_contraction_span_pos(self):
        sl = 50.0
        dc = 112.0
        do = dc + sl
        theta = math.radians(60.0)
        r, phi = nhf.geometry.min_radius_contraction_span_pos(do, dc, theta)
        with self.subTest(state='open'):
            x = r * math.cos(phi)
            y = r * math.sin(phi)
            d = math.sqrt((x - r) ** 2 + y ** 2)
            self.assertAlmostEqual(d, do)
        with self.subTest(state='closed'):
            x = r * math.cos(phi - theta)
            y = r * math.sin(phi - theta)
            d = math.sqrt((x - r) ** 2 + y ** 2)
            self.assertAlmostEqual(d, dc)
    def test_min_tangent_contraction_span_pos(self):
        sl = 50.0
        dc = 112.0
        do = dc + sl
        theta = math.radians(60.0)
        r, phi, rp = nhf.geometry.min_tangent_contraction_span_pos(do, dc, theta)
        with self.subTest(state='open'):
            x = r * math.cos(phi)
            y = r * math.sin(phi)
            d = math.sqrt((x - rp) ** 2 + y ** 2)
            self.assertAlmostEqual(d, do)
        with self.subTest(state='closed'):
            x = r * math.cos(phi - theta)
            y = r * math.sin(phi - theta)
            d = math.sqrt((x - rp) ** 2 + y ** 2)
            self.assertAlmostEqual(d, dc)
    def test_contraction_span_pos_from_radius(self):
        sl = 50.0
        dc = 112.0
        do = dc + sl
        r = 70.0
        theta = math.radians(60.0)
        for smaller in [False, True]:
            with self.subTest(smaller=smaller):
                r, phi, rp = nhf.geometry.contraction_span_pos_from_radius(do, dc, r=r, theta=theta, smaller=smaller)
                with self.subTest(state='open'):
                    x = r * math.cos(phi)
                    y = r * math.sin(phi)
                    d = math.sqrt((x - rp) ** 2 + y ** 2)
                    self.assertAlmostEqual(d, do)
                with self.subTest(state='closed'):
                    x = r * math.cos(phi - theta)
                    y = r * math.sin(phi - theta)
                    d = math.sqrt((x - rp) ** 2 + y ** 2)
                    self.assertAlmostEqual(d, dc)
    #def test_contraction_span_pos_from_radius_2(self):
    #    sl = 40.0
    #    dc = 170.0
    #    do = dc + sl
    #    r = 50.0
    #    theta = math.radians(120.0)
    #    for smaller in [False, True]:
    #        with self.subTest(smaller=smaller):
    #            r, phi, rp = nhf.geometry.contraction_span_pos_from_radius(do, dc, r=r, theta=theta, smaller=smaller)
    #            with self.subTest(state='open'):
    #                x = r * math.cos(phi)
    #                y = r * math.sin(phi)
    #                d = math.sqrt((x - rp) ** 2 + y ** 2)
    #                self.assertAlmostEqual(d, do)
    #            with self.subTest(state='closed'):
    #                x = r * math.cos(phi - theta)
    #                y = r * math.sin(phi - theta)
    #                d = math.sqrt((x - rp) ** 2 + y ** 2)
    #                self.assertAlmostEqual(d, dc)


class TestUtils(unittest.TestCase):

    def test_2d_orientation(self):
        l1 = Cq.Location.from2d(1.2, 0)
        l2 = Cq.Location.from2d(0, 0, 90)
        l3 = l2 * l1
        (x, y), r = l3.to2d()
        self.assertAlmostEqual(x, 0)
        self.assertAlmostEqual(y, 1.2)
        self.assertAlmostEqual(r, 90)

    def test_2d_planar(self):
        l1 = Cq.Location.from2d(1.2, 4.5, 67)
        l2 = Cq.Location.from2d(98, 5.4, 36)
        l3 = Cq.Location.from2d(10, 10, 0)
        l = l3 * l2 * l1
        self.assertTrue(l.is2d())

    def test_tag_point(self):
        """
        A board with 3 holes of unequal sizes. Each hole is marked
        """
        p4x, p4y = 5, 5
        p3x, p3y = 0, 0
        p2x, p2y = -5, 0
        board = (
            Cq.Workplane('XY')
            .box(15, 15, 5)
            .faces("<Z")
            .workplane()
            .pushPoints([(p4x, p4y)])
            .hole(4, depth=2)
            .pushPoints([(p3x, p3y)])
            .hole(3, depth=1.5)
            .pushPoints([(p2x, p2y)])
            .hole(2, depth=1)
        )
        board.moveTo(p4x, p4y).tagPoint("h4")
        board.moveTo(p3x, p3y).tagPoint("h3")
        board.moveTo(p2x, p2y).tagPoint("h2")

        assembly = (
            Cq.Assembly()
            .add(board, name="board", color=color_parent)
            .add(makeSphere(2), name="s4", color=color_child)
            .add(makeSphere(1.5), name="s3", color=color_child)
            .add(makeSphere(1), name="s2", color=color_child)
            .constrain("board?h4", "s4", "Point")
            .constrain("board?h3", "s3", "Point")
            .constrain("board?h2", "s2", "Point")
            .solve()
        )
        self.assertEqual(nhf.checks.pairwise_intersection(assembly), [])
        bbox = assembly.toCompound().BoundingBox()
        self.assertAlmostEqual(bbox.xlen, 15)
        self.assertAlmostEqual(bbox.ylen, 15)
        self.assertAlmostEqual(bbox.zlen, 7)

    def test_tag_plane(self):
        p4x, p4y = 5, 5
        p3x, p3y = 0, 0
        p2x, p2y = -5, 0
        board = (
            Cq.Workplane('XY')
            .box(15, 15, 5)
            .faces("<Z")
            .workplane()
            .pushPoints([(p4x, p4y)])
            .hole(4, depth=2)
            .pushPoints([(p3x, p3y)])
            .hole(3, depth=1.5)
            .pushPoints([(p2x, p2y)])
            .hole(2, depth=1)
        )
        board.moveTo(p4x, p4y).tagPlane("h4")
        board.moveTo(p3x, p3y).tagPlane("h3")
        board.moveTo(p2x, p2y).tagPlane("h2")

        def markedCylOf(r):
            cyl = (
                Cq.Workplane('XY')
                .cylinder(radius=r, height=r)
            )
            cyl.faces("<Z").tag("mate")
            return cyl

        assembly = (
            Cq.Assembly()
            .add(board, name="board", color=color_parent)
            .add(markedCylOf(2), name="c4", color=color_child)
            .add(markedCylOf(1.5), name="c3", color=color_child)
            .add(markedCylOf(1), name="c2", color=color_child)
            .constrain("board?h4", "c4?mate", "Plane", param=0)
            .constrain("board?h3", "c3?mate", "Plane", param=0)
            .constrain("board?h2", "c2?mate", "Plane", param=0)
            .solve()
        )
        self.assertEqual(nhf.checks.pairwise_intersection(assembly), [])
        bbox = assembly.toCompound().BoundingBox()
        self.assertAlmostEqual(bbox.xlen, 15)
        self.assertAlmostEqual(bbox.ylen, 15)
        self.assertAlmostEqual(bbox.zlen, 5)

    # FIXME: Absolute location
    #def test_abs_location(self):
    #    box = Cq.Solid.makeBox(1, 1, 1)
    #    assembly = (
    #        Cq.Assembly()
    #        .add(box, name="b1")
    #        .add(box, name="b2", loc=Cq.Location((0,0,1)))
    #        .add(box, name="b3", loc=Cq.Location((0,0,2)))
    #    )
    #    (x, y, z), _ = assembly.get_abs_location("b2@faces@>Y").toTuple()
    #    self.assertAlmostEqual(x, 0.5)
    #    self.assertAlmostEqual(y, 1)
    #    self.assertAlmostEqual(z, 1.5)
    #    (rx, ry, rz), _ = assembly.get_abs_direction("b2@faces@>Y").toTuple()
    #    self.assertAlmostEqual(rx, 0)
    #    self.assertAlmostEqual(ry, 1)
    #    self.assertAlmostEqual(rz, 0)

    def test_centre_of_mass(self):
        assembly = (
            Cq.Assembly()
            .add(MassBall(mass=3).assembly(), name="s1", loc=Cq.Location((0, 0, 0)))
            .add(MassBall(mass=7).assembly(), name="s2", loc=Cq.Location((0, 0, 10)))
        )
        com = assembly.centre_of_mass()
        self.assertAlmostEqual(com.x, 0)
        self.assertAlmostEqual(com.y, 0)
        self.assertAlmostEqual(com.z, 7)

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