Render helper with block containing Arbre element.

lib/arbre/element.rb

@@ -178,11 +178,20 @@ def method_missing(name, *args, &block)
       elsif assigns && assigns.has_key?(name)
         assigns[name]
       elsif helpers.respond_to?(name)
-        helpers.send(name, *args, &block)
+        helper_capture(name, *args, &block)
       else
         super
       end
     end
 
+    # The helper might have a block that builds Arbre elements
+    # which will be rendered (#to_s) inside ActionView::Base#capture.
+    # We do not want such elements added to self, so we push a dummy
+    # current_arbre_element.
+    def helper_capture(name, *args, &block)
+      s = ""
+      within(Element.new) { s = helpers.send(name, *args, &block) }
+      s.is_a?(Element) ? s.to_s : s
+    end
   end
 end

lib/arbre/rails/template_handler.rb

@@ -1,3 +1,21 @@
+# frozen_string_literal: true
+
+ActionView::Base.class_eval do
+  def capture(*args)
+    value = nil
+    buffer = with_output_buffer { value = yield(*args) }
+
+    # Override to handle Arbre elements inside helper blocks.
+    # See https://github.com/rails/rails/issues/17661
+    # and https://github.com/rails/rails/pull/18024#commitcomment-8975180
+    value = value.to_s if value.is_a?(Arbre::Element)
+
+    if (string = buffer.presence || value) && string.is_a?(String)
+      ERB::Util.html_escape string
+    end
+  end
+end
+
 module Arbre
   module Rails
     class TemplateHandler

spec/rails/integration/rendering_spec.rb

@@ -30,6 +30,10 @@ def render_partial_with_instance_variable
     @my_instance_var = "From Instance Var"
     render "arbre/page_with_arb_partial_and_assignment"
   end
+
+  def render_page_with_helpers
+    render "arbre/page_with_helpers"
+  end
 end
 
 
@@ -80,4 +84,15 @@ def render_partial_with_instance_variable
     expect(body).to have_selector("p", text: "Partial: From Instance Var")
   end
 
+  it "should render a page with helpers" do
+    get "/test/render_page_with_helpers"
+    expect(response).to be_success
+    expect(body).to eq <<EOS
+<span>before h1 link</span>
+<h1><a href="/h1_link_path">h1 link text</a></h1>
+<span>before link_to block</span>
+<a href="/link_path">  <i class=\"link-class\">Link text</i>
+</a><span>at end</span>
+EOS
+  end
 end

spec/rails/templates/arbre/page_with_helpers.arb

@@ -0,0 +1,7 @@
+span { 'before h1 link' }
+h1 { link_to('/h1_link_path') { 'h1 link text' } }
+
+span { 'before link_to block' }
+text_node link_to('/link_path') { i("Link text", class: 'link-class') }
+
+span { 'at end' }