1 files changed, 12 insertions, 4 deletions
diff --git a/actionpack/lib/action_dispatch/testing/integration.rb b/actionpack/lib/action_dispatch/testing/integration.rb
index c300a4ea0d..a9a1576fed 100644
--- a/actionpack/lib/action_dispatch/testing/integration.rb
+++ b/actionpack/lib/action_dispatch/testing/integration.rb
@@ -326,13 +326,21 @@ module ActionDispatch
         @integration_session = Integration::Session.new(app)
       end
 
+      def remove! # :nodoc:
+        @integration_session = nil
+      end
+
       %w(get post patch put head delete cookies assigns
          xml_http_request xhr get_via_redirect post_via_redirect).each do |method|
         define_method(method) do |*args|
           reset! unless integration_session
-          reset_template_assertion
-          # reset the html_document variable, but only for new get/post calls
-          @html_document = nil unless method == 'cookies' || method == 'assigns'
+
+          # reset the html_document variable, except for cookies/assigns calls
+          unless method == 'cookies' || method == 'assigns'
+            @html_document = nil
+            reset_template_assertion
+          end
+
           integration_session.__send__(method, *args).tap do
             copy_session_variables!
           end
@@ -497,7 +505,7 @@ module ActionDispatch
     end
 
     def document_root_element
-      html_document
+      html_document.root
     end
   end
 end