From fbb048ed32066243b823a9b08c0204da6f765ea6 Mon Sep 17 00:00:00 2001
From: Piers Chambers <piers@varyonic.com>
Date: Mon, 6 Feb 2017 15:37:26 -0500
Subject: [PATCH] Render helper with block.

---
 lib/arbre/element.rb                           | 11 ++++++++++-
 lib/arbre/rails/template_handler.rb            | 18 ++++++++++++++++++
 spec/rails/integration/rendering_spec.rb       | 15 +++++++++++++++
 .../templates/arbre/page_with_helpers.arb      |  7 +++++++
 4 files changed, 50 insertions(+), 1 deletion(-)
 create mode 100644 spec/rails/templates/arbre/page_with_helpers.arb

diff --git a/lib/arbre/element.rb b/lib/arbre/element.rb
index 3697f513..3dfde498 100644
--- a/lib/arbre/element.rb
+++ b/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
diff --git a/lib/arbre/rails/template_handler.rb b/lib/arbre/rails/template_handler.rb
index d834629a..b2dda7d0 100644
--- a/lib/arbre/rails/template_handler.rb
+++ b/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
diff --git a/spec/rails/integration/rendering_spec.rb b/spec/rails/integration/rendering_spec.rb
index cdeae1c7..fa676d4d 100644
--- a/spec/rails/integration/rendering_spec.rb
+++ b/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
diff --git a/spec/rails/templates/arbre/page_with_helpers.arb b/spec/rails/templates/arbre/page_with_helpers.arb
new file mode 100644
index 00000000..7a59ad67
--- /dev/null
+++ b/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' }
\ No newline at end of file