If error occurs within a within_frame block, save the offending frame

Gemfile

@@ -7,3 +7,9 @@ gem 'rack', '~> 1.0'
 gem 'rake', '~> 10.0'
 gem 'appraisal', '~> 2.0'
 gem 'aruba', '~> 0.14.0'
+# gem 'debug'
+
+# For using selenium driver
+gem 'selenium-webdriver'
+gem 'webdrivers'
+gem 'puma', '~> 5.6'

lib/capybara-screenshot/capybara.rb

@@ -29,10 +29,18 @@ def screenshot_and_open_image
   end
 
   module SessionScreenshotOverrides
-    def within_window(window_or_handle)
+    def within_window(window)
       super
     rescue Exception
-      Thread.current[:capybara_screenshot_offending_window] = window_or_handle
+      Thread.current[:capybara_screenshot_offending_window] = window
+
+      raise
+    end
+
+    def within_frame(frame)
+      super
+    rescue Exception
+      Thread.current[:capybara_screenshot_offending_frame] = frame
 
       raise
     end

lib/capybara-screenshot/saver.rb

@@ -27,7 +27,7 @@ def initialize(capybara, page, html_save=true, filename_prefix='screenshot')
 
       def save
         current_path do |path|
-          within_offending_window do
+          within_offending_window_or_frame do
             if path.empty?
               warn 'WARN: Screenshot could not be saved. `page.current_path` is empty.'
             else
@@ -145,15 +145,35 @@ def which(cmd)
         nil
       end
 
+      def within_offending_window_or_frame
+        within_offending_window do
+          within_offending_frame do
+            yield
+          end
+        end
+      end
+
       def within_offending_window
         return yield unless Thread.current[:capybara_screenshot_offending_window]
 
         page.within_window(Thread.current[:capybara_screenshot_offending_window]) do
           yield
         end
 
+      ensure
         Thread.current[:capybara_screenshot_offending_window] = nil
       end
+
+      def within_offending_frame
+        return yield unless Thread.current[:capybara_screenshot_offending_frame]
+
+        page.within_frame(Thread.current[:capybara_screenshot_offending_frame]) do
+          yield
+        end
+
+      ensure
+        Thread.current[:capybara_screenshot_offending_frame] = nil
+      end
     end
   end
 end

spec/feature/capybara_spec.rb

@@ -0,0 +1,76 @@
+require 'spec_helper'
+
+describe Capybara::Screenshot::RSpec, 'using with Capybara', type: :aruba do
+  include CommonSetup
+
+  before do
+    Capybara::Screenshot.capybara_tmp_path = expand_path('tmp')
+  end
+
+  def run_failing_case(code, error_message, format=nil)
+    run_case code, format: format
+    if error_message.kind_of?(Regexp)
+      expect(last_command_started.output).to match(error_message)
+    else
+      expect(last_command_started.output).to include(error_message)
+    end
+  end
+
+  def run_case(code, options = {})
+    write_file('spec/test_failure.rb', <<-RUBY)
+      #{ensure_load_paths_valid}
+      require 'rspec'
+      require 'capybara'
+      require 'capybara/rspec'
+      require 'capybara-screenshot'
+      require 'capybara-screenshot/rspec'
+
+      require 'webdrivers'
+      Capybara.register_driver :firefox do |app|
+        options = Selenium::WebDriver::Firefox::Options.new
+        options.args << '--headless'
+        capabilities = Selenium::WebDriver::Remote::Capabilities.firefox({})
+
+        Capybara::Selenium::Driver.new(
+          app,
+          browser: :firefox,
+          capabilities: [capabilities, options],
+        )
+      end
+      Capybara.default_driver = :firefox
+
+      #{setup_test_app}
+      #{code}
+    RUBY
+
+    cmd = cmd_with_format(options[:format])
+    run_simple_with_retry cmd
+
+    expect(last_command_started.output).to match('0 failures') if options[:assert_all_passed]
+  end
+
+  def cmd_with_format(format)
+    "rspec #{"--format #{format} " if format}#{expand_path('spec/test_failure.rb')}"
+  end
+
+  # TODO: within_window did not have tests before
+  it 'saves a screenshot for the correct window for failures ocurring inside within_window'
+
+  it 'saves a screenshot for the correct frame for failures ocurring inside within_frame' do
+    run_failing_case <<-RUBY, %r{Unable to find (visible )?link or button "you'll never find me"}
+      RSpec.describe '/has_frame', type: :feature do
+        it do
+          visit '/has_frame'
+          expect(page.body).to include('This is the has_frame page')
+
+          within_frame 'different_page_frame' do
+            puts page.body
+            expect(page.body).to include('This is a different page')
+            click_on "you'll never find me"
+          end
+        end
+      end
+    RUBY
+    expect('tmp/screenshot.html').to have_file_content(/This is a different page/)
+  end
+end

spec/support/test_app.rb

@@ -10,4 +10,10 @@ class TestApp < Sinatra::Base
   get '/different_page' do
     'This is a different page'
   end
+
+  get '/has_frame' do
+    'This is the has_frame page
+    <iframe src="/different_page" id="different_page_frame">
+    </iframe'
+  end
 end