From bdba2354f7bc72d033d709569902405f86a29948 Mon Sep 17 00:00:00 2001 From: Max Chesterfield Date: Thu, 1 Oct 2026 13:52:36 +1000 Subject: [PATCH] EquipmentTreeBuilder no longer overwrites existing leaves with children. Signed-off-by: Max Chesterfield --- changelog.md | 1 + .../actions/equipment_tree_builder.py | 7 ++- .../actions/test_equipment_tree_builder.py | 54 ++++++++++++++++++- 3 files changed, 59 insertions(+), 3 deletions(-) diff --git a/changelog.md b/changelog.md index ee265bd01..e8fdccd57 100644 --- a/changelog.md +++ b/changelog.md @@ -39,6 +39,7 @@ * **F401**: Removed unused imports across multiple files. * **F403**: Resolved star-import usage by adding per-file-ignores in `pyproject.toml` for modules that intentionally use `from module import *`. * **E402**: Added inline `noqa` comments for intentional module-level imports placed mid-file to avoid circular import issues. +* EquipmentTreeBuilder no longer overwrites existing leaves with children. ### Notes * The lint environment requires `ruff` to be installed. Run `tox -e lint` to check code quality. diff --git a/src/zepben/ewb/services/network/tracing/networktrace/actions/equipment_tree_builder.py b/src/zepben/ewb/services/network/tracing/networktrace/actions/equipment_tree_builder.py index 607feba20..de0b3c0b9 100644 --- a/src/zepben/ewb/services/network/tracing/networktrace/actions/equipment_tree_builder.py +++ b/src/zepben/ewb/services/network/tracing/networktrace/actions/equipment_tree_builder.py @@ -92,9 +92,12 @@ def _apply(self, _: NetworkTraceStep[Any], context: StepContext): self._process_leaf(current_node) def _process_leaf(self, current_node: TreeNode[ConductingEquipment]): - self._leaves.add(current_node) # add this node to _leaves as it has no children + if current_node.children: + self._leaves.discard(current_node) + else: + self._leaves.add(current_node) if current_node.parent: - self._leaves.discard(current_node.parent) # this nodes parent now has a child, it's not a leaf anymore + self._leaves.discard(current_node.parent) # this nodes parent now has a child, it's not a leaf anymore def clear(self): self.roots.clear() diff --git a/test/services/network/tracing/networktrace/actions/test_equipment_tree_builder.py b/test/services/network/tracing/networktrace/actions/test_equipment_tree_builder.py index 4c00601d5..8892fdc24 100644 --- a/test/services/network/tracing/networktrace/actions/test_equipment_tree_builder.py +++ b/test/services/network/tracing/networktrace/actions/test_equipment_tree_builder.py @@ -10,8 +10,9 @@ from services.network.test_data.looping_network import create_looping_network from services.network.tracing.feeder.direction_logger import log_directions -from zepben.ewb import ConductingEquipment, Tracing, NetworkStateOperators, TestNetworkBuilder +from zepben.ewb import ConductingEquipment, Tracing, NetworkStateOperators, TestNetworkBuilder, stop_at_open from zepben.ewb import downstream, NetworkTraceActionType +from zepben.ewb.services.network.tracing.networktrace import tracing from zepben.ewb.services.network.tracing.networktrace.actions.equipment_tree_builder import EquipmentTreeBuilder from zepben.ewb.services.network.tracing.networktrace.actions.tree_node import TreeNode @@ -53,6 +54,57 @@ async def test_equipment_tree_builder_leaves(): assert ce in {l.identified_object for l in tree_builder.leaves} +@pytest.mark.asyncio +async def test_both_sides_of_an_open_switch_are_added_as_leaves_without_adding_the_branching_equipment(): + """ + j0 + |acls1 + | + j4--acls3----j2--acls13-j12 + | | + acls5 acls11 + | | + j6--acls7---b8---acls9--j10 + :return: + """ + + n = ( + TestNetworkBuilder() + .from_junction(num_terminals=1) + .to_acls() + .to_junction(num_terminals=3) + .to_acls() + .to_junction() + .to_acls() + .to_junction() + .to_acls() + .to_breaker(is_open=True, is_normally_open=True) + .to_acls() + .to_junction() + .to_acls() + .to_junction() + .to_acls() + .connect('c13', 'j2', 2, 3) + ).network + + head_junction = n.get('j0') + + await Tracing.set_direction().run(head_junction) + + builder = EquipmentTreeBuilder(calculate_leaves=True) + + await ( + Tracing.network_trace(action_step_type=NetworkTraceActionType.ALL_STEPS) + .add_condition(downstream()) + .add_step_action(builder) + .add_condition(stop_at_open()) + ).run(head_junction) + + assert len(builder.leaves) == 2 + assert {leaf.identified_object for leaf in builder.leaves} == {n.get('b8')} + assert n.get('j2') not in {leaf.identified_object for leaf in builder.leaves} + + @pytest.mark.asyncio async def test_downstream_tree(): n = create_looping_network()