use ancestry in replace_children

app/models/mixins/relationship_mixin.rb

@@ -586,7 +586,7 @@ def add_parent(parent)
   end
 
   def add_children(*child_objs)
-    return child_objs.each { |child| child.with_relationship_type(relationship_type) { child.update!(:parent => self) } } if use_ancestry?
+    return child_objs.flatten.each { |child| child.with_relationship_type(relationship_type) { child.update!(:parent => self) } } if use_ancestry?
 
     options = child_objs.extract_options!
     child_objs = child_objs.flatten
@@ -653,6 +653,14 @@ def replace_children(*child_objs)
     child_objs = child_objs.flatten
     return remove_all_children if child_objs.empty?
 
+    if use_ancestry?
+      to_add = child_objs.collect(&:id) - children.collect(&:id)
+      to_add, to_keep = child_objs.partition { |c| to_add.include?(c.id) }
+      remove_children(children - to_keep)
+      add_children(to_add)
+      return
+    end
+
     # Determine which child relationships should be destroyed, already exist, or should be added
     child_rels = self.child_rels
 
@@ -706,6 +714,8 @@ def remove_all_parents(*args)
   end
 
   def remove_all_children(*args)
+    return children.each { |child| remove_children(child) } if use_ancestry?
+
     # Determine if we are removing all or some children
     options = args.last.kind_of?(Hash) ? args.last : {}
     of_type = Array.wrap(options[:of_type])

spec/models/mixins/relationship_mixin_spec.rb

@@ -356,6 +356,36 @@ def self.name
         service2.save!
         expect(service6.with_relationship_type("custom") { service6.child_ids }).to eq([service2.id])
       end
+
+      context "#replace_children" do
+        it "with one child to keep and one to add" do
+          service = ancestry_class.create
+          service1.with_relationship_type("custom") { service1.replace_children(service, service4) }
+          expect(service1.with_relationship_type("custom") { service1.reload.children }).to contain_exactly(service, service4)
+        end
+
+        it "with all to remove" do
+          service1.with_relationship_type("custom") { service1.replace_children }
+          expect(service1.with_relationship_type("custom") { service1.replace_children }).to eq([])
+        end
+
+        it "with one child to keep" do
+          service1.with_relationship_type("custom") { service1.replace_children(service4) }
+          expect(service1.with_relationship_type("custom") { service1.reload.children }).to contain_exactly(service4)
+        end
+
+        it "with one child to remove" do
+          service = ancestry_class.create
+          service1.with_relationship_type("custom") { service1.replace_children(service) }
+          expect(service1.with_relationship_type("custom") { service1.reload.children }).to eq([service])
+        end
+
+        it "with no children to start adds them all" do
+          service = ancestry_class.create
+          service.with_relationship_type("custom") { service.replace_children(service4) }
+          expect(service.with_relationship_type("custom") { service.reload.children }).to eq([service4])
+        end
+      end
     end
 
     describe "#add_parent" do