CelerisLab/tests/unit/test_pending_edits.py

"""ObjectManager pending edit lifecycle — stage_add, stage_remove, has_pending_edit, clear_pending_edits.

No GPU required."""

import unittest

from CelerisLab.body.manager import ObjectManager
from CelerisLab.body.objects import SimObject
from CelerisLab.body.geometry.circle import CircleGeometry


def _make_circle_obj(cx: float, cy: float, radius: float,
                     is_sensor: bool = False) -> SimObject:
    """Create a minimal SimObject with CircleGeometry."""
    geom = CircleGeometry(cx, cy, radius)
    return SimObject(
        obj_id=-1,
        geometry=geom,
        center=(cx, cy),
        radius=radius,
        is_sensor=is_sensor,
    )


class TestPendingEditLifecycle(unittest.TestCase):
    """Test the pending edit state machine on ObjectManager."""

    def setUp(self):
        # ObjectManager requires nx, ny, nz, nq, cfg.  Use a minimal cfg stub.
        self.cfg = _StubCfg(dim=2)
        self.mgr = ObjectManager(nx=64, ny=32, nz=1, nq=9, cfg=self.cfg)

    # -- stage_add -----------------------------------------------------------

    def test_stage_add_sets_edit_active(self):
        obj = _make_circle_obj(30, 16, 5)
        self.assertFalse(self.mgr.has_pending_edit())
        self.mgr.stage_add(obj)
        self.assertTrue(self.mgr.has_pending_edit())

    def test_stage_add_does_not_change_formal_count(self):
        obj = _make_circle_obj(30, 16, 5)
        self.assertEqual(self.mgr.count, 0)
        self.mgr.stage_add(obj)
        self.assertEqual(self.mgr.count, 0)

    def test_stage_add_multiple(self):
        for i in range(3):
            self.mgr.stage_add(_make_circle_obj(10 + i * 10, 16, 3))
        self.assertTrue(self.mgr.has_pending_edit())

    # -- stage_remove --------------------------------------------------------

    def test_stage_remove_valid_id(self):
        body_id = self.mgr.add(_make_circle_obj(30, 16, 5))
        self.mgr.stage_remove(body_id)
        self.assertTrue(self.mgr.has_pending_edit())

    def test_stage_remove_invalid_id_raises(self):
        with self.assertRaises(IndexError):
            self.mgr.stage_remove(999)

    def test_stage_remove_does_not_change_formal_count(self):
        body_id = self.mgr.add(_make_circle_obj(30, 16, 5))
        self.assertEqual(self.mgr.count, 1)
        self.mgr.stage_remove(body_id)
        self.assertEqual(self.mgr.count, 1)

    # -- has_pending_edit ----------------------------------------------------

    def test_has_pending_edit_false_initially(self):
        self.assertFalse(self.mgr.has_pending_edit())

    def test_has_pending_edit_false_after_clear(self):
        self.mgr.stage_add(_make_circle_obj(30, 16, 5))
        self.mgr.clear_pending_edits()
        self.assertFalse(self.mgr.has_pending_edit())

    def test_has_pending_edit_false_after_empty_stage(self):
        # If we add and then remove the same pending add, edit is still
        # "active" (edit_active=True) but has no pending content.
        # has_pending_edit should return False.
        obj = _make_circle_obj(30, 16, 5)
        self.mgr.stage_add(obj)
        # Manually clear pending_add to simulate an empty edit window
        self.mgr._pending_add.clear()
        self.assertFalse(self.mgr.has_pending_edit())

    # -- clear_pending_edits -------------------------------------------------

    def test_clear_resets_all_pending(self):
        body_id = self.mgr.add(_make_circle_obj(30, 16, 5))
        self.mgr.stage_add(_make_circle_obj(40, 16, 3))
        self.mgr.stage_remove(body_id)
        self.assertTrue(self.mgr.has_pending_edit())

        self.mgr.clear_pending_edits()
        self.assertFalse(self.mgr.has_pending_edit())
        self.assertEqual(len(self.mgr._pending_add), 0)
        self.assertEqual(len(self.mgr._pending_remove), 0)
        self.assertFalse(self.mgr._edit_active)

    def test_clear_preserves_formal_registry(self):
        body_id = self.mgr.add(_make_circle_obj(30, 16, 5))
        self.mgr.stage_remove(body_id)
        self.mgr.clear_pending_edits()
        # Formal object should still be there
        self.assertEqual(self.mgr.count, 1)
        self.assertEqual(self.mgr.get(body_id).obj_id, body_id)

    # -- Combination: add + remove -------------------------------------------

    def test_stage_add_and_remove_together(self):
        body_id = self.mgr.add(_make_circle_obj(30, 16, 5))
        self.mgr.stage_add(_make_circle_obj(40, 16, 3))
        self.mgr.stage_remove(body_id)
        self.assertTrue(self.mgr.has_pending_edit())
        # Formal count unchanged
        self.assertEqual(self.mgr.count, 1)

    # -- Formal add (pre-initialize path) ------------------------------------

    def test_formal_add_still_works(self):
        """The existing add() path must remain functional."""
        obj = _make_circle_obj(30, 16, 5)
        body_id = self.mgr.add(obj)
        self.assertEqual(body_id, 0)
        self.assertEqual(self.mgr.count, 1)

    def test_formal_add_and_pending_coexist(self):
        """Formal add + pending stage should not interfere."""
        self.mgr.add(_make_circle_obj(30, 16, 5))
        self.mgr.stage_add(_make_circle_obj(40, 16, 3))
        self.assertEqual(self.mgr.count, 1)
        self.assertTrue(self.mgr.has_pending_edit())


class _StubCfg:
    """Minimal LBMConfig-like stub for ObjectManager construction."""

    def __init__(self, dim: int = 2):
        self.dim = dim
        self.is_d3q19 = (dim == 3)


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