Return to the first window after yield

lib/capybara/driver/webkit/browser.rb

@@ -82,7 +82,11 @@ def set_skip_image_loading(skip_image_loading)
     end
 
     def window_focus(handle=nil)
-      command("WindowFocus")
+      if handle
+        command("WindowFocus", handle)
+      else
+        command("WindowFocus")
+      end
     end
 
     def command(name, *args)

spec/driver_spec.rb

@@ -1560,6 +1560,7 @@ def which_for(character)
               <script type="text/javascript">
                 window.open('http://#{request.host_with_port}/test?#{request.query_string}');
               </script>
+              <p>bananas</p>
             </html>
           HTML
         else
@@ -1575,16 +1576,22 @@ def which_for(character)
 
     it "has the expected text in the new window" do
       subject.visit("/new_window")
-      subject.within_window(nil) do
+      subject.within_window(true) do
         subject.find("//p").first.text.should == "finished"
       end
     end
 
     it "waits for the new window to load" do
       subject.visit("/new_window?sleep=1")
-      subject.within_window(nil) do
+      subject.within_window(true) do
         lambda { Timeout::timeout(1) { subject.find("//p") } }.should_not raise_error(Timeout::Error)
       end
     end
+
+    it "switches back to the original window" do
+      subject.visit("/new_window")
+      subject.within_window(true) { }
+      subject.find("//p").first.text.should == "bananas"
+    end
   end
 end

src/WebPageManager.cpp

@@ -17,6 +17,10 @@ void WebPageManager::append(WebPage *value) {
   webPages->append(value);
 }
 
+WebPage *WebPageManager::first() {
+  return (WebPage*) webPages->first();
+}
+
 WebPage *WebPageManager::last() {
   return (WebPage*) webPages->last();
 }

src/WebPageManager.h

@@ -5,6 +5,7 @@ class WebPageManager {
   public:
     static WebPageManager *getInstance();
     void append(WebPage *value);
+    WebPage *first();
     WebPage *last();
 
   private:

src/WindowFocus.cpp

@@ -7,12 +7,16 @@ WindowFocus::WindowFocus(WebPage *page, QStringList &arguments, QObject *parent)
 }
 
 void WindowFocus::start() {
-  WebPage *webPage = WebPageManager::getInstance()->last();
-  if (webPage) {
-    emit windowChanged(webPage);
-    success();
-  } else {
-    windowNotFound();
+  WebPageManager *manager = WebPageManager::getInstance();
+
+  switch(arguments().length()) {
+    case 1:
+      emit windowChanged(manager->last());
+      success();
+      break;
+    default:
+      emit windowChanged(manager->first());
+      success();
   }
 }