From 468bec886f15a2de7703155d3a88035eda1dd4bc Mon Sep 17 00:00:00 2001
From: Matthew Horan <matt@matthoran.com>
Date: Tue, 27 Mar 2012 13:14:01 -0400
Subject: [PATCH] Return to the first window after yield

---
 lib/capybara/driver/webkit/browser.rb |  6 +++++-
 spec/driver_spec.rb                   | 11 +++++++++--
 src/WebPageManager.cpp                |  4 ++++
 src/WebPageManager.h                  |  1 +
 src/WindowFocus.cpp                   | 16 ++++++++++------
 5 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/lib/capybara/driver/webkit/browser.rb b/lib/capybara/driver/webkit/browser.rb
index e160165..5be287f 100644
--- a/lib/capybara/driver/webkit/browser.rb
+++ b/lib/capybara/driver/webkit/browser.rb
@@ -82,7 +82,11 @@ class Capybara::Driver::Webkit
     end
 
     def window_focus(handle=nil)
-      command("WindowFocus")
+      if handle
+        command("WindowFocus", handle)
+      else
+        command("WindowFocus")
+      end
     end
 
     def command(name, *args)
diff --git a/spec/driver_spec.rb b/spec/driver_spec.rb
index 5900743..847d647 100644
--- a/spec/driver_spec.rb
+++ b/spec/driver_spec.rb
@@ -1560,6 +1560,7 @@ describe Capybara::Driver::Webkit do
               <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 @@ describe Capybara::Driver::Webkit do
 
     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
diff --git a/src/WebPageManager.cpp b/src/WebPageManager.cpp
index 5e90c6b..e729ee9 100644
--- a/src/WebPageManager.cpp
+++ b/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();
 }
diff --git a/src/WebPageManager.h b/src/WebPageManager.h
index d5e7ee7..a1b7613 100644
--- a/src/WebPageManager.h
+++ b/src/WebPageManager.h
@@ -5,6 +5,7 @@ class WebPageManager {
   public:
     static WebPageManager *getInstance();
     void append(WebPage *value);
+    WebPage *first();
     WebPage *last();
 
   private:
diff --git a/src/WindowFocus.cpp b/src/WindowFocus.cpp
index 1a8cb4b..3406cb1 100644
--- a/src/WindowFocus.cpp
+++ b/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();
   }
 }