From 0a9d6302c3cc706e9aa025e31813b3e7d1af0a61 Mon Sep 17 00:00:00 2001
From: lulalala <mark@goodlife.tw>
Date: Wed, 13 Jul 2022 21:41:39 +0800
Subject: [PATCH] Take ownership when children nodes are assigned

---
 lib/oga/xml/document.rb        |  1 +
 lib/oga/xml/node.rb            |  1 +
 lib/oga/xml/node_set.rb        | 14 +++++++++-----
 spec/oga/xml/generator_spec.rb | 11 +++++++++++
 4 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/lib/oga/xml/document.rb b/lib/oga/xml/document.rb
index 1b62664..a689e50 100644
--- a/lib/oga/xml/document.rb
+++ b/lib/oga/xml/document.rb
@@ -47,6 +47,7 @@ module Oga
       def children=(nodes)
         if nodes.is_a?(NodeSet)
           nodes.owner = self
+          nodes.take_ownership_on_nodes
           @children = nodes
         else
           @children = NodeSet.new(nodes, self)
diff --git a/lib/oga/xml/node.rb b/lib/oga/xml/node.rb
index 8330c2a..6e2326c 100644
--- a/lib/oga/xml/node.rb
+++ b/lib/oga/xml/node.rb
@@ -50,6 +50,7 @@ module Oga
       def children=(nodes)
         if nodes.is_a?(NodeSet)
           nodes.owner = self
+          nodes.take_ownership_on_nodes
           @children = nodes
         else
           @children = NodeSet.new(nodes, self)
diff --git a/lib/oga/xml/node_set.rb b/lib/oga/xml/node_set.rb
index feb508e..712b97c 100644
--- a/lib/oga/xml/node_set.rb
+++ b/lib/oga/xml/node_set.rb
@@ -42,11 +42,7 @@ module Oga
         @owner    = owner
         @existing = {}
 
-        @nodes.each_with_index do |node, index|
-          mark_existing(node)
-
-          take_ownership(node, index) if @owner
-        end
+        take_ownership_on_nodes
       end
 
       # Yields the supplied block for every node.
@@ -289,6 +285,14 @@ module Oga
         "NodeSet(#{values})"
       end
 
+      def take_ownership_on_nodes
+        @nodes.each_with_index do |node, index|
+          mark_existing(node)
+
+          take_ownership(node, index) if @owner
+        end
+      end
+
       private
 
       # Takes ownership of the given node. This only occurs when the current
diff --git a/spec/oga/xml/generator_spec.rb b/spec/oga/xml/generator_spec.rb
index 9202058..b392b69 100644
--- a/spec/oga/xml/generator_spec.rb
+++ b/spec/oga/xml/generator_spec.rb
@@ -168,5 +168,16 @@ describe Oga::XML::Generator do
         end
       end
     end
+
+    describe 'using an Element with replaced children' do
+      it 'returns a string' do
+        element = Oga::XML::Element.new(name: 'foo')
+        element.children = Oga::XML::Parser.new('<bar></bar>').parse.children
+
+        output = described_class.new(element).to_xml
+
+        expect(output).to eq('<foo><bar /></foo>')
+      end
+    end
   end
 end