Fix bug in restore_ancestry_integrity!, added test to verify we didn'…

…t remove the ancestry from all nodes.

lib/ancestry/class_methods.rb

@@ -107,33 +107,36 @@ def check_ancestry_integrity! options = {}
     # Integrity restoration
     def restore_ancestry_integrity!
       parents = {}
-      # For each node ...
-      self.base_class.find_each do |node|
-        # ... set its ancestry to nil if invalid
-        if node.errors[node.class.ancestry_column].blank?
-          node.without_ancestry_callbacks do
-            node.update_attribute node.ancestry_column, nil
+      # Wrap the whole thing in a transaction ...
+      self.base_class.transaction do
+        # For each node ...
+        self.base_class.find_each do |node|
+          # ... set its ancestry to nil if invalid
+          if !node.valid? and !node.errors[node.class.ancestry_column].blank?
+            node.without_ancestry_callbacks do
+              node.update_attribute node.ancestry_column, nil
+            end
           end
-        end
-        # ... save parent of this node in parents array if it exists
-        parents[node.id] = node.parent_id if exists? node.parent_id
+          # ... save parent of this node in parents array if it exists
+          parents[node.id] = node.parent_id if exists? node.parent_id
 
-        # Reset parent id in array to nil if it introduces a cycle
-        parent = parents[node.id]
-        until parent.nil? || parent == node.id
-          parent = parents[parent]
-        end
-        parents[node.id] = nil if parent == node.id 
-      end
-      # For each node ...
-      self.base_class.find_each do |node|
-        # ... rebuild ancestry from parents array
-        ancestry, parent = nil, parents[node.id]
-        until parent.nil?
-          ancestry, parent = if ancestry.nil? then parent else "#{parent}/#{ancestry}" end, parents[parent]
+          # Reset parent id in array to nil if it introduces a cycle
+          parent = parents[node.id]
+          until parent.nil? || parent == node.id
+            parent = parents[parent]
+          end
+          parents[node.id] = nil if parent == node.id 
         end
-        node.without_ancestry_callbacks do
-          node.update_attribute node.ancestry_column, ancestry
+        # For each node ...
+        self.base_class.find_each do |node|
+          # ... rebuild ancestry from parents array
+          ancestry, parent = nil, parents[node.id]
+          until parent.nil?
+            ancestry, parent = if ancestry.nil? then parent else "#{parent}/#{ancestry}" end, parents[parent]
+          end
+          node.without_ancestry_callbacks do
+            node.update_attribute node.ancestry_column, ancestry
+          end
         end
       end
     end

test/has_ancestry_test.rb

@@ -379,23 +379,26 @@ def assert_integrity_restoration model
     assert_nothing_raised do
       model.check_ancestry_integrity!
     end
+    assert model.all.any? {|node| node.ancestry.present? }, "Expected some nodes not to be roots"
+    assert_equal model.count, model.roots.collect {|node| node.descendants.count + 1 }.sum
   end
 
   def test_integrity_restoration
+    width, depth = 3, 3
     # Check that integrity is restored for invalid format for ancestry column
-    AncestryTestDatabase.with_model :width => 3, :depth => 3 do |model, roots|
+    AncestryTestDatabase.with_model :width => width, :depth => depth do |model, roots|
       roots.first.first.update_attribute model.ancestry_column, 'invalid_ancestry'
       assert_integrity_restoration model
     end
 
     # Check that integrity is restored for non-existent ancestor
-    AncestryTestDatabase.with_model :width => 3, :depth => 3 do |model, roots|
+    AncestryTestDatabase.with_model :width => width, :depth => depth do |model, roots|
       roots.first.first.update_attribute model.ancestry_column, 35
       assert_integrity_restoration model
     end
 
     # Check that integrity is restored for cyclic ancestry
-    AncestryTestDatabase.with_model :width => 3, :depth => 3 do |model, roots|
+    AncestryTestDatabase.with_model :width => width, :depth => depth do |model, roots|
       node = roots.first.first
       node.update_attribute model.ancestry_column, node.id
       assert_integrity_restoration model