Add option on how to handle orphans when arranging nodes

lib/ancestry/class_methods.rb

@@ -40,16 +40,29 @@ def arrange(options = {})
     # arranges array of nodes to a hierarchical hash
     #
     # @param nodes [Array[Node]] nodes to be arranged
+    # @param orphan_strategy [Symbol]  :rootify or :destroy (default: :rootify)
     # @returns Hash{Node => {Node => {}, Node => {}}}
     # If a node's parent is not included, the node will be included as if it is a top level node
-    def arrange_nodes(nodes)
+    def arrange_nodes(nodes, orphan_strategy: :rootify)
       node_ids = Set.new(nodes.map(&:id))
       index = Hash.new { |h, k| h[k] = {} }
 
       nodes.each_with_object({}) do |node, arranged|
-        children = index[node.id]
-        index[node.parent_id][node] = children
-        arranged[node] = children unless node_ids.include?(node.parent_id)
+        index[node.parent_id][node] = children = index[node.id]
+        if node.parent_id.nil?
+          arranged[node] = children
+        elsif !node_ids.include?(node.parent_id)
+          case orphan_strategy
+          when :destroy
+             # All children are destroyed as well (default)
+          when :adopt
+            raise ArgumentError, "Not Implemented"
+          when :rootify
+            arranged[node] = children
+          when :restrict
+            raise Ancestry::AncestryException, I18n.t("ancestry.cannot_delete_descendants")
+          end
+        end
       end
     end
 

test/concerns/arrangement_test.rb

@@ -234,4 +234,61 @@ def test_arrange_partial
       assert_equal([n5.id, n3.id], model.arrange_nodes([n5, n3]).keys.map(&:id))
     end
   end
+
+  def test_arrange_nodes_orphan_strategy_rootify
+    AncestryTestDatabase.with_model do |model|
+      # - n1
+      #   - n2
+      #     - n3
+      #   - n4
+      n1 = model.create!
+      n2 = model.create!(parent: n1)
+      n3 = model.create!(parent: n2)
+      n4 = model.create!(parent: n1)
+
+      # n2 and n4's parent (n1) is missing, so they become roots
+      result = model.arrange_nodes([n2, n3, n4], orphan_strategy: :rootify)
+      assert_equal({n2 => {n3 => {}}, n4 => {}}, result)
+    end
+  end
+
+  def test_arrange_nodes_orphan_strategy_destroy
+    AncestryTestDatabase.with_model do |model|
+      # - n1
+      #   - n2
+      #     - n3
+      #   - n4
+      n1 = model.create!
+      n2 = model.create!(parent: n1)
+      n3 = model.create!(parent: n2)
+      n4 = model.create!(parent: n1)
+
+      # n2 and n4's parent (n1) is missing, so they and their children are dropped
+      result = model.arrange_nodes([n2, n3, n4], orphan_strategy: :destroy)
+      assert_equal({}, result)
+    end
+  end
+
+  def test_arrange_nodes_orphan_strategy_restrict
+    AncestryTestDatabase.with_model do |model|
+      n1 = model.create!
+      n2 = model.create!(parent: n1)
+
+      # n2's parent (n1) is missing, so restrict raises
+      assert_raises(Ancestry::AncestryException) do
+        model.arrange_nodes([n2], orphan_strategy: :restrict)
+      end
+    end
+  end
+
+  def test_arrange_nodes_orphan_strategy_rootify_is_default
+    AncestryTestDatabase.with_model do |model|
+      n1 = model.create!
+      n2 = model.create!(parent: n1)
+
+      # default behavior: orphans become roots
+      result = model.arrange_nodes([n2])
+      assert_equal({n2 => {}}, result)
+    end
+  end
 end