Bug 1259055 - Use window management class for handling of new windows. draft
authorHenrik Skupin <mail@hskupin.info>
Thu, 03 Nov 2016 17:46:32 +0100
changeset 433362 a18a737c15d2a02b0ba2fc6d1007b3e389c8b0ec
parent 433361 97f83ced176dcb4f097a61ce258e0b0f2b1b1ac9
child 535874 815bf56a04c4bd0d2bb3b6cfceb8daca22c427f3
push id34562
push userbmo:hskupin@gmail.com
push dateThu, 03 Nov 2016 17:00:59 +0000
bugs1259055
milestone52.0a1
Bug 1259055 - Use window management class for handling of new windows. MozReview-Commit-ID: 4dwHLFEmr4k
testing/marionette/harness/marionette/runner/mixins/window_manager.py
testing/marionette/harness/marionette/tests/unit/test_anonymous_content.py
testing/marionette/harness/marionette/tests/unit/test_checkbox_chrome.py
testing/marionette/harness/marionette/tests/unit/test_chrome.py
testing/marionette/harness/marionette/tests/unit/test_elementsize_chrome.py
testing/marionette/harness/marionette/tests/unit/test_execute_script.py
testing/marionette/harness/marionette/tests/unit/test_import_script.py
testing/marionette/harness/marionette/tests/unit/test_key_actions.py
testing/marionette/harness/marionette/tests/unit/test_pagesource_chrome.py
testing/marionette/harness/marionette/tests/unit/test_screenshot.py
testing/marionette/harness/marionette/tests/unit/test_switch_frame_chrome.py
testing/marionette/harness/marionette/tests/unit/test_text_chrome.py
testing/marionette/harness/marionette/tests/unit/test_window_handles.py
testing/marionette/harness/marionette/tests/unit/test_window_management.py
testing/marionette/harness/marionette/tests/unit/test_window_switching.py
testing/marionette/harness/marionette/www/windowHandles.html
--- a/testing/marionette/harness/marionette/runner/mixins/window_manager.py
+++ b/testing/marionette/harness/marionette/runner/mixins/window_manager.py
@@ -9,25 +9,30 @@ from marionette_driver import By, Wait
 
 class WindowManagerMixin(object):
 
     _menu_item_new_tab = (By.ID, "menu_newNavigatorTab")
 
     def setUp(self):
         super(WindowManagerMixin, self).setUp()
 
+        self.start_window = self.marionette.current_chrome_window_handle
+        self.start_windows = self.marionette.chrome_window_handles
+
         self.start_tab = self.marionette.current_window_handle
         self.start_tabs = self.marionette.window_handles
 
     def tearDown(self):
-        try:
-            if len(self.marionette.window_handles) != len(self.start_tabs):
-                raise Exception("Not all tabs as opened by the test have been closed")
-        finally:
-            super(WindowManagerMixin, self).tearDown()
+        if len(self.marionette.chrome_window_handles) != len(self.start_windows):
+            raise Exception("Not all windows as opened by the test have been closed")
+
+        if len(self.marionette.window_handles) != len(self.start_tabs):
+            raise Exception("Not all tabs as opened by the test have been closed")
+
+        super(WindowManagerMixin, self).tearDown()
 
     def close_all_tabs(self):
         current_window_handles = self.marionette.window_handles
 
         # If the start tab is not present anymore, use the next one of the list
         if self.start_tab not in current_window_handles:
             self.start_tab = current_window_handles[0]
 
@@ -39,16 +44,37 @@ class WindowManagerMixin(object):
             # Bug 1311350 - close() doesn't wait for tab to be closed.
             Wait(self.marionette).until(
                 lambda mn: handle not in mn.window_handles,
                 message="Failed to close tab with handle {}".format(handle)
             )
 
         self.marionette.switch_to_window(self.start_tab)
 
+    def close_all_windows(self):
+        current_chrome_window_handles = self.marionette.chrome_window_handles
+
+        # If the start window is not present anymore, use the next one of the list
+        if self.start_window not in current_chrome_window_handles:
+            self.start_window = current_chrome_window_handles[0]
+        current_chrome_window_handles.remove(self.start_window)
+
+        with self.marionette.using_context("chrome"):
+            for handle in current_chrome_window_handles:
+                self.marionette.switch_to_window(handle)
+                self.marionette.close_chrome_window()
+
+                # Bug 1311350 - close_chrome_window() doesn't wait for window to be closed.
+                Wait(self.marionette).until(
+                    lambda mn: handle not in mn.chrome_window_handles,
+                    message="Failed to close window with handle {}".format(handle)
+                )
+
+            self.marionette.switch_to_window(self.start_window)
+
     def open_tab(self, trigger="menu"):
         current_tabs = self.marionette.window_handles
 
         try:
             if callable(trigger):
                 trigger()
             elif trigger == 'menu':
                 with self.marionette.using_context("chrome"):
@@ -60,8 +86,44 @@ class WindowManagerMixin(object):
             Wait(self.marionette).until(
                 lambda mn: len(mn.window_handles) == len(current_tabs) + 1,
                 message="No new tab has been opened"
             )
 
             [new_tab] = list(set(self.marionette.window_handles) - set(current_tabs))
 
             return new_tab
+
+    def open_window(self, trigger=None):
+        current_windows = self.marionette.chrome_window_handles
+
+        def loaded(handle):
+            with self.marionette.using_context('chrome'):
+                return self.marionette.execute_script("""
+                  Components.utils.import("resource://gre/modules/Services.jsm");
+
+                  let win = Services.wm.getOuterWindowWithId(Number(arguments[0]));
+                  return win.document.readyState == "complete";
+                """, script_args=[handle])
+
+        try:
+            if callable(trigger):
+                trigger()
+            else:
+                with self.marionette.using_context("chrome"):
+                    self.marionette.execute_script("window.open();")
+        except Exception:
+            exc, val, tb = sys.exc_info()
+            raise exc, 'Failed to trigger opening a new window: {}'.format(val), tb
+        else:
+            Wait(self.marionette).until(
+                lambda mn: len(mn.chrome_window_handles) == len(current_windows) + 1,
+                message="No new window has been opened"
+            )
+
+            [new_window] = list(set(self.marionette.chrome_window_handles) - set(current_windows))
+
+            # Before continuing ensure the window has been completed loading
+            Wait(self.marionette).until(
+                lambda _: loaded(new_window),
+                message="Window with handle '{}'' did not finish loading".format(new_window))
+
+            return new_window
--- a/testing/marionette/harness/marionette/tests/unit/test_anonymous_content.py
+++ b/testing/marionette/harness/marionette/tests/unit/test_anonymous_content.py
@@ -1,45 +1,51 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-from marionette import MarionetteTestCase
+from unittest import skip
 
+from marionette import MarionetteTestCase, WindowManagerMixin
 from marionette_driver.errors import NoSuchElementException
 from marionette_driver.expected import element_present
 from marionette_driver.marionette import HTMLElement
 from marionette_driver.wait import Wait
 from marionette_driver.by import By
 
 
-class TestAnonymousContent(MarionetteTestCase):
+class TestAnonymousContent(WindowManagerMixin, MarionetteTestCase):
+
     def setUp(self):
-        MarionetteTestCase.setUp(self)
+        super(TestAnonymousContent, self).setUp()
         self.marionette.set_context("chrome")
-        self.win = self.marionette.current_chrome_window_handle
-        self.marionette.execute_script("window.open('chrome://marionette/content/test_anonymous_content.xul', 'foo', 'chrome,centerscreen');")
-        self.marionette.switch_to_window('foo')
-        self.assertNotEqual(self.win, self.marionette.current_chrome_window_handle)
+
+        def open_window_with_js():
+            self.marionette.execute_script("""
+              window.open('chrome://marionette/content/test_anonymous_content.xul',
+                          'foo', 'chrome,centerscreen');
+            """)
+
+        new_window = self.open_window(trigger=open_window_with_js)
+        self.marionette.switch_to_window(new_window)
+        self.assertNotEqual(self.marionette.current_chrome_window_handle, self.start_window)
 
     def tearDown(self):
-        self.assertNotEqual(self.win, self.marionette.current_chrome_window_handle)
-        self.marionette.close_chrome_window()
-        self.marionette.switch_to_window(self.win)
-        MarionetteTestCase.tearDown(self)
+        self.close_all_windows()
 
     def test_switch_to_anonymous_frame(self):
         self.marionette.find_element(By.ID, "testAnonymousContentBox")
         anon_browser_el = self.marionette.find_element(By.ID, "browser")
         self.assertTrue("test_anonymous_content.xul" in self.marionette.get_url())
         self.marionette.switch_to_frame(anon_browser_el)
         self.assertTrue("test.xul" in self.marionette.get_url())
         self.marionette.find_element(By.ID, "testXulBox")
         self.assertRaises(NoSuchElementException, self.marionette.find_element, By.ID, "testAnonymousContentBox")
 
+    @skip("Bug 1311657 - Opened chrome window cannot be closed after call to switch_to_frame")
     def test_switch_to_anonymous_iframe(self):
         self.marionette.find_element(By.ID, "testAnonymousContentBox")
         el = self.marionette.find_element(By.ID, "container2")
         anon_iframe_el = el.find_element(By.ANON_ATTRIBUTE, {"anonid": "iframe"})
         self.marionette.switch_to_frame(anon_iframe_el)
         self.assertTrue("test.xul" in self.marionette.get_url())
         self.marionette.find_element(By.ID, "testXulBox")
         self.assertRaises(NoSuchElementException, self.marionette.find_element, By.ID,
--- a/testing/marionette/harness/marionette/tests/unit/test_checkbox_chrome.py
+++ b/testing/marionette/harness/marionette/tests/unit/test_checkbox_chrome.py
@@ -1,29 +1,35 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-from marionette import MarionetteTestCase
+from marionette import MarionetteTestCase, WindowManagerMixin
 from marionette_driver.by import By
 
 
-class TestSelectedChrome(MarionetteTestCase):
+class TestSelectedChrome(WindowManagerMixin, MarionetteTestCase):
+
     def setUp(self):
-        MarionetteTestCase.setUp(self)
+        super(TestSelectedChrome, self).setUp()
+
         self.marionette.set_context("chrome")
-        self.win = self.marionette.current_chrome_window_handle
-        self.marionette.execute_script("window.open('chrome://marionette/content/test.xul', '_blank', 'chrome,centerscreen');")
+
+        def open_window_with_js():
+            self.marionette.execute_script("""
+              window.open('chrome://marionette/content/test.xul',
+                          '_blank', 'chrome,centerscreen');
+            """)
+
+        new_window = self.open_window(trigger=open_window_with_js)
+        self.marionette.switch_to_window(new_window)
 
     def tearDown(self):
-        self.marionette.close_chrome_window()
-        self.marionette.switch_to_window(self.win)
-        MarionetteTestCase.tearDown(self)
+        try:
+            self.close_all_windows()
+        finally:
+            super(TestSelectedChrome, self).tearDown()
 
     def test_selected(self):
-        wins = self.marionette.chrome_window_handles
-        wins.remove(self.win)
-        newWin = wins.pop()
-        self.marionette.switch_to_window(newWin)
         box = self.marionette.find_element(By.ID, "testBox")
         self.assertFalse(box.is_selected())
         self.assertFalse(self.marionette.execute_script("arguments[0].checked = true;", [box]))
         self.assertTrue(box.is_selected())
--- a/testing/marionette/harness/marionette/tests/unit/test_chrome.py
+++ b/testing/marionette/harness/marionette/tests/unit/test_chrome.py
@@ -8,28 +8,43 @@
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
 #Unless required by applicable law or agreed to in writing, software
 #distributed under the License is distributed on an "AS IS" BASIS,
 #WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #See the License for the specific language governing permissions and
 #limitations under the License.
 
+from marionette import MarionetteTestCase, WindowManagerMixin
 from marionette_driver import By
-from marionette_driver.errors import NoSuchElementException
-from marionette import MarionetteTestCase
 
 
-class ChromeTests(MarionetteTestCase):
+class ChromeTests(WindowManagerMixin, MarionetteTestCase):
+
+    def setUp(self):
+        super(ChromeTests, self).setUp()
+
+        self.marionette.set_context('chrome')
+
+    def tearDown(self):
+        self.close_all_windows()
+        super(ChromeTests, self).tearDown()
 
     def test_hang_until_timeout(self):
-        with self.marionette.using_context('chrome'):
-            start_handle = self.marionette.current_chrome_window_handle
+        def open_with_menu():
             menu = self.marionette.find_element(By.ID, 'aboutName')
             menu.click()
-            handles = self.marionette.chrome_window_handles
-            handles.remove(start_handle)
-            self.marionette.switch_to_window(handles[0])
-            self.assertRaises(NoSuchElementException, self.marionette.find_element, By.ID, 'dek')
+
+        new_window = self.open_window(trigger=open_with_menu)
+        self.marionette.switch_to_window(new_window)
 
-            # Clean up the window
-            self.marionette.close()
-            self.marionette.switch_to_window(start_handle)
+        try:
+            try:
+                # Raise an exception type which should not be thrown by Marionette
+                # while running this test. Otherwise it would mask eg. IOError as
+                # thrown for a socket timeout.
+                raise NotImplementedError('Exception should not cause a hang when '
+                                          'closing the chrome window')
+            finally:
+                self.marionette.close_chrome_window()
+                self.marionette.switch_to_window(self.start_window)
+        except NotImplementedError:
+            pass
--- a/testing/marionette/harness/marionette/tests/unit/test_elementsize_chrome.py
+++ b/testing/marionette/harness/marionette/tests/unit/test_elementsize_chrome.py
@@ -1,33 +1,34 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-from marionette import MarionetteTestCase
+
+from marionette import MarionetteTestCase, WindowManagerMixin
 from marionette_driver.by import By
 
 
-class TestElementSizeChrome(MarionetteTestCase):
+class TestElementSizeChrome(WindowManagerMixin, MarionetteTestCase):
+
     def setUp(self):
-        MarionetteTestCase.setUp(self)
+        super(TestElementSizeChrome, self).setUp()
+
         self.marionette.set_context("chrome")
-        self.win = self.marionette.current_chrome_window_handle
-        self.marionette.execute_script(
-            "window.open('chrome://marionette/content/test2.xul', 'foo', 'chrome,centerscreen');")
-        self.marionette.switch_to_window('foo')
-        self.assertNotEqual(self.win, self.marionette.current_chrome_window_handle)
+
+        def open_window_with_js():
+            self.marionette.execute_script("""
+              window.open('chrome://marionette/content/test2.xul',
+                          'foo', 'chrome,centerscreen');
+            """)
+
+        new_window = self.open_window(trigger=open_window_with_js)
+        self.marionette.switch_to_window(new_window)
 
     def tearDown(self):
-        self.assertNotEqual(self.win, self.marionette.current_chrome_window_handle)
-        self.marionette.close_chrome_window()
-        self.marionette.switch_to_window(self.win)
-        MarionetteTestCase.tearDown(self)
+        self.close_all_windows()
+        super(TestElementSizeChrome, self).tearDown()
 
     def testShouldReturnTheSizeOfAnInput(self):
-        wins = self.marionette.chrome_window_handles
-        wins.remove(self.win)
-        newWin = wins.pop()
-        self.marionette.switch_to_window(newWin)
         shrinko = self.marionette.find_element(By.ID, 'textInput')
         size = shrinko.rect
         self.assertTrue(size['width'] > 0)
         self.assertTrue(size['height'] > 0)
--- a/testing/marionette/harness/marionette/tests/unit/test_execute_script.py
+++ b/testing/marionette/harness/marionette/tests/unit/test_execute_script.py
@@ -1,18 +1,18 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
+import os
 import urllib
-import os
 
+from marionette import MarionetteTestCase, WindowManagerMixin
 from marionette_driver import By, errors
 from marionette_driver.marionette import HTMLElement
-from marionette import MarionetteTestCase
 
 
 def inline(doc):
     return "data:text/html;charset=utf-8,{}".format(urllib.quote(doc))
 
 
 elements = inline("<p>foo</p> <p>bar</p>")
 
@@ -233,30 +233,32 @@ class TestExecuteContent(MarionetteTestC
         self.assertTrue(
             send("return typeof window.wrappedJSObject == 'undefined'"))
 
     def test_no_callback(self):
         self.assertTrue(self.marionette.execute_script(
             "return typeof arguments[0] == 'undefined'"))
 
 
-class TestExecuteChrome(TestExecuteContent):
+class TestExecuteChrome(WindowManagerMixin, TestExecuteContent):
+
     def setUp(self):
-        TestExecuteContent.setUp(self)
-        self.win = self.marionette.current_window_handle
+        super(TestExecuteChrome, self).setUp()
+
         self.marionette.set_context("chrome")
-        self.marionette.execute_script(
-            "window.open('chrome://marionette/content/test.xul', 'xul', 'chrome')")
-        self.marionette.switch_to_window("xul")
-        self.assertNotEqual(self.win, self.marionette.current_window_handle)
+
+        def open_window_with_js():
+            self.marionette.execute_script(
+                "window.open('chrome://marionette/content/test.xul', 'xul', 'chrome');")
+
+        new_window = self.open_window(trigger=open_window_with_js)
+        self.marionette.switch_to_window(new_window)
 
     def tearDown(self):
-        self.marionette.execute_script("window.close()")
-        self.marionette.switch_to_window(self.win)
-        self.assertEqual(self.win, self.marionette.current_window_handle)
+        self.close_all_windows()
         super(TestExecuteChrome, self).tearDown()
 
     def test_permission(self):
         self.assertEqual(1, self.marionette.execute_script(
             "var c = Components.classes; return 1;"))
 
     def test_unmarshal_element_collection(self):
         expected = self.marionette.find_elements(By.TAG_NAME, "textbox")
--- a/testing/marionette/harness/marionette/tests/unit/test_import_script.py
+++ b/testing/marionette/harness/marionette/tests/unit/test_import_script.py
@@ -1,29 +1,30 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 import os
 
-from marionette.marionette_test import MarionetteTestCase, skip_if_chrome
+from marionette import MarionetteTestCase, WindowManagerMixin, skip_if_chrome
 from marionette_driver.errors import JavascriptException
 from marionette_driver.by import By
 
 
-class TestImportScriptContent(MarionetteTestCase):
+class TestImportScriptContent(WindowManagerMixin, MarionetteTestCase):
     contexts = set(["chrome", "content"])
 
     script_file = os.path.abspath(
         os.path.join(__file__, os.path.pardir, "importscript.js"))
     another_script_file = os.path.abspath(
         os.path.join(__file__, os.path.pardir, "importanotherscript.js"))
 
     def setUp(self):
-        MarionetteTestCase.setUp(self)
+        super(TestImportScriptContent, self).setUp()
+
         for context in self.contexts:
             with self.marionette.using_context(context):
                 self.marionette.clear_imported_scripts()
         self.reset_context()
 
     def reset_context(self):
         self.marionette.set_context("content")
 
@@ -101,27 +102,25 @@ class TestImportScriptContent(Marionette
         self.marionette.import_script(self.another_script_file)
         self.assert_defined("testFunc")
         self.assert_defined("testAnotherFunc")
 
     @skip_if_chrome
     def test_imports_apply_globally(self):
         self.marionette.navigate(
             self.marionette.absolute_url("test_windows.html"))
-        original_window = self.marionette.current_chrome_window_handle
-        self.marionette.find_element(By.LINK_TEXT, "Open new window").click()
 
-        windows = set(self.marionette.chrome_window_handles)
-        print "windows={}".format(windows)
-        new_window = windows.difference([original_window]).pop()
+        def open_window_with_link():
+            self.marionette.find_element(By.LINK_TEXT, "Open new window").click()
+
+        new_window = self.open_window(trigger=open_window_with_link)
         self.marionette.switch_to_window(new_window)
 
         self.marionette.import_script(self.script_file)
         self.marionette.close_chrome_window()
 
-        print "switching to original window: {}".format(original_window)
-        self.marionette.switch_to_window(original_window)
+        self.marionette.switch_to_window(self.start_window)
         self.assert_defined("testFunc")
 
 
 class TestImportScriptChrome(TestImportScriptContent):
     def reset_context(self):
         self.marionette.set_context("chrome")
--- a/testing/marionette/harness/marionette/tests/unit/test_key_actions.py
+++ b/testing/marionette/harness/marionette/tests/unit/test_key_actions.py
@@ -1,21 +1,22 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-from marionette import MarionetteTestCase
+from marionette import MarionetteTestCase, WindowManagerMixin
 from marionette_driver.keys import Keys
 from marionette_driver.marionette import Actions
 from marionette_driver.by import By
 
 
-class TestKeyActions(MarionetteTestCase):
+class TestKeyActions(WindowManagerMixin, MarionetteTestCase):
+
     def setUp(self):
-        MarionetteTestCase.setUp(self)
+        super(TestKeyActions, self).setUp()
         if self.marionette.session_capabilities["platformName"] == "darwin":
             self.mod_key = Keys.META
         else:
             self.mod_key = Keys.CONTROL
         test_html = self.marionette.absolute_url("javascriptPage.html")
         self.marionette.navigate(test_html)
         self.reporter_element = self.marionette.find_element(By.ID, "keyReporter")
         self.reporter_element.click()
@@ -64,27 +65,24 @@ class TestKeyActions(MarionetteTestCase)
         (self.key_action.key_down(self.mod_key)
                         .key_down("a")
                         .wait(.5)
                         .key_down("x")
                         .perform())
         self.assertEqual(self.key_reporter_value, "")
 
     def test_open_in_new_window_shortcut(self):
-        start_win = self.marionette.current_chrome_window_handle
-        el = self.marionette.find_element(By.ID, "updatediv")
-        # Ensure that the element is in the current view port because press() doesn't
-        # handle that inside the action chain (bug 1295538).
-        self.marionette.execute_script('arguments[0].scrollIntoView()', script_args=[el])
-        (self.key_action.key_down(Keys.SHIFT)
-                        .press(el)
-                        .release()
-                        .key_up(Keys.SHIFT)
-                        .perform())
-        self.wait_for_condition(
-            lambda mn: len(self.marionette.chrome_window_handles) == 2)
-        chrome_window_handles = self.marionette.chrome_window_handles
-        chrome_window_handles.remove(start_win)
-        [new_win] = chrome_window_handles
-        self.marionette.switch_to_window(new_win)
+        def open_window_with_action():
+            el = self.marionette.find_element(By.ID, "updatediv")
+            # Ensure that the element is in the current view port because press() doesn't
+            # handle that inside the action chain (bug 1295538).
+            self.marionette.execute_script('arguments[0].scrollIntoView()', script_args=[el])
+            (self.key_action.key_down(Keys.SHIFT)
+                            .press(el)
+                            .release()
+                            .key_up(Keys.SHIFT)
+                            .perform())
+
+        new_window = self.open_window(trigger=open_window_with_action)
+        self.marionette.switch_to_window(new_window)
         self.marionette.close_chrome_window()
-        self.marionette.switch_to_window(start_win)
+        self.marionette.switch_to_window(self.start_window)
         self.assertEqual(self.key_reporter_value, "")
--- a/testing/marionette/harness/marionette/tests/unit/test_pagesource_chrome.py
+++ b/testing/marionette/harness/marionette/tests/unit/test_pagesource_chrome.py
@@ -1,29 +1,29 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-from marionette import MarionetteTestCase
+from marionette import MarionetteTestCase, WindowManagerMixin
 
 
-class TestPageSourceChrome(MarionetteTestCase):
+class TestPageSourceChrome(WindowManagerMixin, MarionetteTestCase):
+
     def setUp(self):
-        MarionetteTestCase.setUp(self)
+        super(TestPageSourceChrome, self).setUp()
         self.marionette.set_context("chrome")
-        self.win = self.marionette.current_chrome_window_handle
-        self.marionette.execute_script("window.open('chrome://marionette/content/test.xul', 'foo', 'chrome,centerscreen');")
-        self.marionette.switch_to_window('foo')
-        self.assertNotEqual(self.win, self.marionette.current_chrome_window_handle)
+
+        def open_with_js():
+            self.marionette.execute_script("""
+              window.open('chrome://marionette/content/test.xul',
+                          'foo', 'chrome,centerscreen');
+            """)
+
+        new_window = self.open_window(open_with_js)
+        self.marionette.switch_to_window(new_window)
 
     def tearDown(self):
-        self.assertNotEqual(self.win, self.marionette.current_chrome_window_handle)
-        self.marionette.close_chrome_window()
-        self.marionette.switch_to_window(self.win)
-        MarionetteTestCase.tearDown(self)
+        self.close_all_windows()
+        super(TestPageSourceChrome, self).tearDown()
 
     def testShouldReturnXULDetails(self):
-        wins = self.marionette.chrome_window_handles
-        wins.remove(self.win)
-        newWin = wins.pop()
-        self.marionette.switch_to_window(newWin)
-        source  = self.marionette.page_source
+        source = self.marionette.page_source
         self.assertTrue('<textbox id="textInput"' in source)
--- a/testing/marionette/harness/marionette/tests/unit/test_screenshot.py
+++ b/testing/marionette/harness/marionette/tests/unit/test_screenshot.py
@@ -1,22 +1,21 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 import base64
 import hashlib
 import imghdr
 import struct
-import time
 import urllib
 
 from unittest import skip
 
-from marionette import MarionetteTestCase
+from marionette import MarionetteTestCase, WindowManagerMixin
 from marionette_driver.by import By
 
 
 def inline(doc, mime="text/html;charset=utf-8"):
     return "data:{0},{1}".format(mime, urllib.quote(doc))
 
 
 ELEMENT = "iVBORw0KGgoAAAANSUhEUgAAADIAAAAyCAYAAAAeP4ixAAAAVklEQVRoge3PMQ0AMAzAsPJHVWYbjEWTj/zx7O75oXk9AAISD6QWSC2QWiC1QGqB1AKpBVILpBZILZBaILVAaoHUAqkFUgukFkgtkFogtUBqgdT6BnIBMKa1DtYxhPkAAAAASUVORK5CYII="
@@ -40,37 +39,37 @@ class ScreenCaptureTestCase(MarionetteTe
 
     def get_image_dimensions(self, string):
         self.assert_png(string)
         image = base64.decodestring(string)
         width, height = struct.unpack(">LL", image[16:24])
         return int(width), int(height)
 
 
-class Chrome(ScreenCaptureTestCase):
+class TestScreenCaptureChrome(WindowManagerMixin, ScreenCaptureTestCase):
     @property
     def primary_window_dimensions(self):
         current_window = self.marionette.current_chrome_window_handle
-        self.marionette.switch_to_window(self.original_window)
+        self.marionette.switch_to_window(self.start_window)
         with self.marionette.using_context("chrome"):
             rv = tuple(self.marionette.execute_script("""
                 let el = document.documentElement;
                 let rect = el.getBoundingClientRect();
                 return [rect.width, rect.height];
                 """))
         self.marionette.switch_to_window(current_window)
         return rv
 
     def setUp(self):
-        ScreenCaptureTestCase.setUp(self)
+        super(TestScreenCaptureChrome, self).setUp()
         self.marionette.set_context("chrome")
-        self.original_window = self.marionette.current_chrome_window_handle
 
     def tearDown(self):
-        self.marionette.switch_to_window(self.original_window)
+        self.close_all_windows()
+        super(TestScreenCaptureChrome, self).tearDown()
 
     # A full chrome window screenshot is not the outer dimensions of
     # the window, but instead the bounding box of the <window> inside
     # <browser>.
     def test_window(self):
         ss = self.marionette.screenshot()
         self.assert_png(ss)
         self.assertEqual(self.primary_window_dimensions,
@@ -81,37 +80,32 @@ class Chrome(ScreenCaptureTestCase):
             content = self.marionette.screenshot()
         chrome = self.marionette.screenshot()
         self.assertNotEqual(content, chrome)
 
     # This tests that GeckoDriver#takeScreenshot uses
     # currentContext.document.documentElement instead of looking for a
     # <window> element, which does not exist for all windows.
     def test_secondary_windows(self):
-        self.assertEqual(len(self.marionette.chrome_window_handles), 1)
-        ss = self.marionette.screenshot()
-        self.marionette.execute_script("""
-            window.open('chrome://marionette/content/test_dialog.xul', 'foo',
-                        'dialog,height=200,width=300');
-            """)
-        self.assertEqual(len(self.marionette.chrome_window_handles), 2)
-        # TODO: Bug 1288769 prevents us from switching to the dialog via its name
-        # self.marionette.switch_to_window('foo')
-        new_win = None
-        for win in self.marionette.chrome_window_handles:
-            if win != self.original_window:
-                new_win = win
-                break
-        self.marionette.switch_to_window(new_win)
+        def open_window_with_js():
+            self.marionette.execute_script("""
+                window.open('chrome://marionette/content/test_dialog.xul', 'foo',
+                            'dialog,height=200,width=300');
+                """)
+
+        new_window = self.open_window(open_window_with_js)
+        self.marionette.switch_to_window(new_window)
+
         ss = self.marionette.screenshot()
         size = self.get_image_dimensions(ss)
         self.assert_png(ss)
         self.assertNotEqual(self.primary_window_dimensions, size)
 
         self.marionette.close_chrome_window()
+        self.marionette.switch_to_window(self.start_window)
 
 
 class Content(ScreenCaptureTestCase):
     @property
     def body_scroll_dimensions(self):
         return tuple(self.marionette.execute_script(
             "return [document.body.scrollWidth, document.body.scrollHeight]"))
 
--- a/testing/marionette/harness/marionette/tests/unit/test_switch_frame_chrome.py
+++ b/testing/marionette/harness/marionette/tests/unit/test_switch_frame_chrome.py
@@ -1,31 +1,39 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-from marionette import MarionetteTestCase
+from unittest import skip
+
+from marionette import MarionetteTestCase, WindowManagerMixin
 from marionette_driver.errors import JavascriptException
 
 
-class TestSwitchFrameChrome(MarionetteTestCase):
+class TestSwitchFrameChrome(WindowManagerMixin, MarionetteTestCase):
+
     def setUp(self):
-        MarionetteTestCase.setUp(self)
+        super(TestSwitchFrameChrome, self).setUp()
         self.marionette.set_context("chrome")
-        self.win = self.marionette.current_window_handle
-        self.marionette.execute_script("window.open('chrome://marionette/content/test.xul', 'foo', 'chrome,centerscreen');")
-        self.marionette.switch_to_window('foo')
-        self.assertNotEqual(self.win, self.marionette.current_window_handle)
+
+        def open_window_with_js():
+            self.marionette.execute_script("""
+              window.open('chrome://marionette/content/test.xul',
+                          'foo', 'chrome,centerscreen');
+            """)
+
+        new_window = self.open_window(trigger=open_window_with_js)
+        self.marionette.switch_to_window(new_window)
+        self.assertNotEqual(self.start_window, self.marionette.current_chrome_window_handle)
 
     def tearDown(self):
-        self.assertNotEqual(self.win, self.marionette.current_window_handle)
-        self.marionette.execute_script("window.close();")
-        self.marionette.switch_to_window(self.win)
-        MarionetteTestCase.tearDown(self)
+        self.close_all_windows()
+        super(TestSwitchFrameChrome, self).tearDown()
 
+    @skip("Bug 1311657 - Opened chrome window cannot be closed after call to switch_to_frame(0)")
     def test_switch_simple(self):
         self.assertIn("test.xul", self.marionette.get_url(), "Initial navigation has failed")
         self.marionette.switch_to_frame(0)
         self.assertIn("test2.xul", self.marionette.get_url(),"Switching by index failed")
         self.marionette.switch_to_frame()
         self.assertEqual(None, self.marionette.get_active_frame(), "Switiching by null failed")
         self.assertIn("test.xul", self.marionette.get_url(), "Switching by null failed")
         self.marionette.switch_to_frame("iframe")
@@ -35,16 +43,17 @@ class TestSwitchFrameChrome(MarionetteTe
         self.marionette.switch_to_frame("iframename")
         self.assertIn("test2.xul", self.marionette.get_url(), "Switching by name failed")
         iframe_element = self.marionette.get_active_frame()
         self.marionette.switch_to_frame()
         self.assertIn("test.xul", self.marionette.get_url(), "Switching by null failed")
         self.marionette.switch_to_frame(iframe_element)
         self.assertIn("test2.xul", self.marionette.get_url(), "Switching by element failed")
 
+    @skip("Bug 1311657 - Opened chrome window cannot be closed after call to switch_to_frame(0)")
     def test_stack_trace(self):
         self.assertIn("test.xul", self.marionette.get_url(), "Initial navigation has failed")
         self.marionette.switch_to_frame(0)
         self.assertRaises(JavascriptException, self.marionette.execute_async_script, "foo();")
         try:
             self.marionette.execute_async_script("foo();")
         except JavascriptException as e:
             self.assertIn("foo", e.message)
--- a/testing/marionette/harness/marionette/tests/unit/test_text_chrome.py
+++ b/testing/marionette/harness/marionette/tests/unit/test_text_chrome.py
@@ -1,52 +1,44 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-from marionette import MarionetteTestCase
+from marionette import MarionetteTestCase, WindowManagerMixin
 from marionette_driver.by import By
 
 
 ''' Disabled in bug 896043 and when working on Chrome code re-enable for bug 896046
-class TestTextChrome(MarionetteTestCase):
+class TestTextChrome(WindowManagerMixin, MarionetteTestCase):
+
     def setUp(self):
-        MarionetteTestCase.setUp(self)
+        super(TestTextChrome, self).setUp()
         self.marionette.set_context("chrome")
-        self.win = self.marionette.current_chrome_window_handle
-        self.marionette.execute_script("window.open('chrome://marionette/content/test.xul', 'foo', 'chrome,centerscreen');")
-        self.marionette.switch_to_window('foo')
-        self.assertNotEqual(self.win, self.marionette.current_chrome_window_handle)
+
+        def open_window_with_js():
+            self.marionette.execute_script("""
+              window.open('chrome://marionette/content/test.xul',
+                          'foo', 'chrome,centerscreen');
+            """)
+
+        new_window = self.open_window(trigger=open_window_with_js)
+        self.marionette.switch_to_window(new_window)
 
     def tearDown(self):
-        self.assertNotEqual(self.win, self.marionette.current_chrome_window_handle)
-        self.marionette.close_chrome_window()
-        self.marionette.switch_to_window(self.win)
-        MarionetteTestCase.tearDown(self)
+        self.close_all_windows()
+        super(TestTextChrome, self).tearDown()
 
     def test_getText(self):
-        wins = self.marionette.chrome_window_handles
-        wins.remove(self.win)
-        newWin = wins.pop()
-        self.marionette.switch_to_window(newWin)
         box = self.marionette.find_element(By.ID, "textInput")
         self.assertEqual("test", box.text)
 
     def test_clearText(self):
-        wins = self.marionette.chrome_window_handles
-        wins.remove(self.win)
-        newWin = wins.pop()
-        self.marionette.switch_to_window(newWin)
         box = self.marionette.find_element(By.ID, "textInput")
         self.assertEqual("test", box.text)
         box.clear()
         self.assertEqual("", box.text)
 
     def test_sendKeys(self):
-        wins = self.marionette.chrome_window_handles
-        wins.remove(self.win)
-        newWin = wins.pop()
-        self.marionette.switch_to_window(newWin)
         box = self.marionette.find_element(By.ID, "textInput")
         self.assertEqual("test", box.text)
         box.send_keys("at")
         self.assertEqual("attest", box.text)
-'''
\ No newline at end of file
+'''
--- a/testing/marionette/harness/marionette/tests/unit/test_window_handles.py
+++ b/testing/marionette/harness/marionette/tests/unit/test_window_handles.py
@@ -4,21 +4,27 @@
 
 from marionette import MarionetteTestCase, WindowManagerMixin
 from marionette_driver.keys import Keys
 from marionette_driver import By, Wait
 
 
 class TestWindowHandles(WindowManagerMixin, MarionetteTestCase):
 
+    def setUp(self):
+        super(TestWindowHandles, self).setUp()
+
+        self.test_page = self.marionette.absolute_url("windowHandles.html")
+        self.marionette.navigate(self.test_page)
+
     def tearDown(self):
-        try:
-            self.close_all_tabs()
-        finally:
-            super(TestWindowHandles, self).tearDown()
+        self.close_all_windows()
+        self.close_all_tabs()
+
+        super(TestWindowHandles, self).tearDown()
 
     def test_new_tab_window_handles(self):
         keys = []
         if self.marionette.session_capabilities['platformName'] == 'darwin':
             keys.append(Keys.META)
         else:
             keys.append(Keys.CONTROL)
         keys.append('t')
@@ -56,151 +62,121 @@ class TestWindowHandles(WindowManagerMix
                 menu_new_tab = self.marionette.find_element(By.ID, 'menu_newNavigatorTab')
                 menu_new_tab.click()
 
         new_tab = self.open_tab(trigger=open_with_menu)
         self.assertEqual(self.marionette.current_window_handle, self.start_tab)
 
         # We still have the default tab set as our window handle. This
         # get_url command should be sent immediately, and not be forever-queued.
-        self.assertEqual(self.marionette.get_url(), "about:blank")
+        self.assertEqual(self.marionette.get_url(), self.test_page)
 
         self.marionette.switch_to_window(new_tab)
         self.marionette.close()
         self.marionette.switch_to_window(self.start_tab)
 
     def test_link_opened_tab_window_handles(self):
-        tab_testpage = self.marionette.absolute_url("windowHandles.html")
-        self.marionette.navigate(tab_testpage)
-
         def open_with_link():
             link = self.marionette.find_element(By.ID, "new-tab")
             link.click()
 
         new_tab = self.open_tab(trigger=open_with_link)
         self.assertEqual(self.marionette.current_window_handle, self.start_tab)
 
         self.marionette.switch_to_window(new_tab)
         self.assertEqual(self.marionette.get_url(), "about:blank")
         self.assertEqual(self.marionette.title, "")
 
         self.marionette.switch_to_window(self.start_tab)
 
-        self.assertEqual(self.marionette.get_url(), tab_testpage)
+        self.assertEqual(self.marionette.get_url(), self.test_page)
         self.assertEqual(self.marionette.title, "Marionette New Tab Link")
 
         self.marionette.close()
         self.marionette.switch_to_window(new_tab)
         self.assertEqual(self.marionette.get_url(), "about:blank")
 
     def test_chrome_windows(self):
-        opener_page = self.marionette.absolute_url("windowHandles.html")
-
-        self.marionette.navigate(opener_page)
-
-        start_win = self.marionette.current_chrome_window_handle
-
-        self.marionette.find_element(By.ID, "new-window").click()
+        # We open a chrome window but are actually interested in the new tab.
+        new_window = self.open_window(
+            trigger=lambda: self.marionette.find_element(By.ID, "new-window").click())
+        with self.marionette.using_context("chrome"):
+            self.marionette.switch_to_window(new_window)
 
-        self.assertEqual(len(self.marionette.chrome_window_handles), 2)
-        self.assertEqual(len(self.marionette.window_handles), 2)
-        windows = self.marionette.chrome_window_handles
-        windows.remove(start_win)
-        dest_win = windows.pop()
-        self.marionette.switch_to_window(dest_win)
+        # Check that the new tab is available and wait until it has been loaded.
+        self.assertEqual(len(self.marionette.window_handles), len(self.start_tabs) + 1)
+        new_tab = self.marionette.current_window_handle
+        Wait(self.marionette).until(lambda _: self.marionette.get_url() == self.test_page,
+                                    message="Test page hasn't been loaded for newly opened tab")
 
-        self.marionette.navigate(opener_page)
-        new_tab_link = self.marionette.find_element(By.ID, "new-tab")
-
+        link_new_tab = self.marionette.find_element(By.ID, "new-tab")
         for i in range(3):
-            new_tab_link.click()
-            self.marionette.switch_to_window(dest_win)
-            self.wait_for_condition(lambda mn: len(mn.window_handles) == i + 3)
-            self.assertEqual(len(self.marionette.chrome_window_handles), 2)
+            self.open_tab(trigger=lambda: link_new_tab.click())
+            self.marionette.switch_to_window(new_tab)
+            # No more chrome windows should be opened
+            self.assertEqual(len(self.marionette.chrome_window_handles),
+                             len(self.start_windows) + 1)
 
         self.marionette.close_chrome_window()
-        self.assertEqual(len(self.marionette.chrome_window_handles), 1)
-        self.assertEqual(len(self.marionette.window_handles), 1)
-        self.marionette.switch_to_window(start_win)
+        self.marionette.switch_to_window(self.start_window)
 
     def test_chrome_window_handles_with_scopes(self):
-        start_tab = self.marionette.current_window_handle
-        start_chrome_window = self.marionette.current_chrome_window_handle
-        start_chrome_windows = self.marionette.chrome_window_handles
-
         # Ensure that we work in chrome scope so we don't have any limitations
         with self.marionette.using_context("chrome"):
             # Open a browser and a non-browser (about window) chrome window
-            self.marionette.execute_script("window.open()")
-            self.marionette.find_element(By.ID, "aboutName").click()
-
-            self.wait_for_condition(lambda mn: len(mn.chrome_window_handles) ==
-                                    len(start_chrome_windows) + 2)
+            self.open_window(
+                trigger=lambda: self.marionette.execute_script("window.open();"))
+            self.open_window(
+                trigger=lambda: self.marionette.find_element(By.ID, "aboutName").click())
 
             handles_in_chrome_scope = self.marionette.chrome_window_handles
             with self.marionette.using_context("content"):
                 self.assertEqual(self.marionette.chrome_window_handles,
                                  handles_in_chrome_scope)
 
-        for handle in handles_in_chrome_scope:
-            if handle != start_chrome_window:
-                self.marionette.switch_to_window(handle)
-                self.marionette.close_chrome_window()
-
-        self.marionette.switch_to_window(start_tab)
-
-        self.assertEqual(len(self.marionette.chrome_window_handles), 1)
-        self.assertEqual(len(self.marionette.window_handles), 1)
-        self.assertEqual(self.marionette.current_window_handle, start_tab)
-
     def test_tab_and_window_handles(self):
-        start_chrome_window = self.marionette.current_chrome_window_handle
-
-        tab_open_page = self.marionette.absolute_url("windowHandles.html")
         window_open_page = self.marionette.absolute_url("test_windows.html")
         results_page = self.marionette.absolute_url("resultPage.html")
 
         # Open a new tab and switch to it.
         def open_tab_with_link():
-            self.marionette.navigate(tab_open_page)
             link = self.marionette.find_element(By.ID, "new-tab")
             link.click()
 
         second_tab = self.open_tab(trigger=open_tab_with_link)
         self.assertEqual(len(self.marionette.chrome_window_handles), 1)
-        self.assertEqual(self.marionette.current_chrome_window_handle, start_chrome_window)
+        self.assertEqual(self.marionette.current_chrome_window_handle, self.start_window)
 
         self.marionette.switch_to_window(second_tab)
         self.assertEqual(self.marionette.get_url(), "about:blank")
 
         # Open a new window from the new tab and only care about the second new tab
-        def open_win_with_link():
-            self.marionette.navigate(window_open_page)
+        def open_window_with_link():
             link = self.marionette.find_element(By.LINK_TEXT, "Open new window")
             link.click()
 
-        third_tab = self.open_tab(trigger=open_win_with_link)
-        self.assertEqual(len(self.marionette.chrome_window_handles), 2)
-        self.assertEqual(self.marionette.current_chrome_window_handle, start_chrome_window)
+        # We open a new window but are actually interested in the new tab
+        self.marionette.navigate(window_open_page)
+        third_tab = self.open_tab(trigger=open_window_with_link)
+        self.assertEqual(self.marionette.current_chrome_window_handle, self.start_window)
 
         # Check that the new tab has the correct page loaded
         self.marionette.switch_to_window(third_tab)
         self.assertEqual(self.marionette.get_url(), results_page)
 
-        self.assertEqual(len(self.marionette.chrome_window_handles), 2)
-        self.assertNotEqual(self.marionette.current_chrome_window_handle, start_chrome_window)
+        self.assertNotEqual(self.marionette.current_chrome_window_handle, self.start_window)
 
         # Return to our original tab and close it.
         self.marionette.switch_to_window(self.start_tab)
         self.marionette.close()
         self.assertEquals(len(self.marionette.window_handles), 2)
 
         # Close the opened window and carry on in our second tab.
         self.marionette.switch_to_window(third_tab)
         self.marionette.close()
         self.assertEqual(len(self.marionette.window_handles), 1)
 
         self.marionette.switch_to_window(second_tab)
         self.assertEqual(self.marionette.get_url(), results_page)
         self.marionette.navigate("about:blank")
 
         self.assertEqual(len(self.marionette.chrome_window_handles), 1)
-        self.assertEqual(self.marionette.current_chrome_window_handle, start_chrome_window)
+        self.assertEqual(self.marionette.current_chrome_window_handle, self.start_window)
--- a/testing/marionette/harness/marionette/tests/unit/test_window_management.py
+++ b/testing/marionette/harness/marionette/tests/unit/test_window_management.py
@@ -1,116 +1,57 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-import time
-from marionette import MarionetteTestCase
-from marionette_driver import By, Wait
+from marionette import MarionetteTestCase, WindowManagerMixin
+from marionette_driver import By
 
 
-class TestSwitchWindow(MarionetteTestCase):
-    def open_new_window(self):
-        with self.marionette.using_context("chrome"):
-            print '**** %s' % self.marionette.execute_async_script("""
-              var ww = Components.classes["@mozilla.org/embedcomp/window-watcher;1"].
-                                  getService(Components.interfaces.nsIWindowWatcher);
-              var win = ww.openWindow(null, "chrome://browser/content/browser.xul",
-                                      "testWin", null, null);
-              win.addEventListener("load", function() {
-                marionetteScriptFinished(true);
-              }, {once: true});
-              dump(win.readyState);
-              if (win.readyState == 'loaded')
-                marionetteScriptFinished(true);
-            """)
+class TestSwitchWindow(WindowManagerMixin, MarionetteTestCase):
 
-    def close_new_window(self):
+    def setUp(self):
+        super(TestSwitchWindow, self).setUp()
         self.marionette.set_context("chrome")
-        self.marionette.execute_script("""
-          var ww = Components.classes["@mozilla.org/embedcomp/window-watcher;1"].
-                              getService(Components.interfaces.nsIWindowWatcher);
-          var win = ww.getWindowByName("testWin", null);
-          if (win != null)
-            win.close();
-        """)
-        self.marionette.set_context("content")
+
+    def tearDown(self):
+        self.close_all_windows()
+        super(TestSwitchWindow, self).tearDown()
 
     def test_windows(self):
-        orig_win = self.marionette.current_chrome_window_handle
-        orig_available = self.marionette.chrome_window_handles
-        self.open_new_window()
-        # assert we're still in the original window
-        self.assertEqual(self.marionette.current_chrome_window_handle, orig_win)
-        # assert we can find the new window
-        Wait(self.marionette).until(
-            lambda _: len(self.marionette.chrome_window_handles) == len(orig_available) + 1,
-            message="The new window has not been opened.")
-        # assert that our window is there
-        now_available = self.marionette.chrome_window_handles
-        self.assertIn(orig_win, now_available)
-        new_win = None
-        for win in now_available:
-            if win != orig_win:
-                new_win = win
-        # switch to another window
-        self.marionette.switch_to_window(new_win)
-        self.assertEqual(self.marionette.current_chrome_window_handle, new_win)
-        self.assertNotEqual(self.marionette.current_chrome_window_handle, orig_win)
-        # switch back
-        self.marionette.switch_to_window(orig_win)
-        self.assertEqual(self.marionette.current_chrome_window_handle, orig_win)
-        self.close_new_window()
-        self.assertNotIn(new_win, self.marionette.chrome_window_handles)
-        self.assertEqual(self.marionette.current_chrome_window_handle, orig_win)
-        self.assertEqual(len(self.marionette.chrome_window_handles), len(orig_available))
+        def open_browser_with_js():
+            self.marionette.execute_script(" window.open(); ")
+
+        new_window = self.open_window(trigger=open_browser_with_js)
+        self.assertEqual(self.marionette.current_chrome_window_handle, self.start_window)
+
+        # switch to the other window
+        self.marionette.switch_to_window(new_window)
+        self.assertEqual(self.marionette.current_chrome_window_handle, new_window)
+        self.assertNotEqual(self.marionette.current_chrome_window_handle, self.start_window)
 
-    def testShouldLoadAWindowAndThenCloseIt(self):
-        test_html = self.marionette.absolute_url("test_windows.html")
-        self.marionette.navigate(test_html)
-        current = self.marionette.current_chrome_window_handle
+        # switch back and close original window
+        self.marionette.switch_to_window(self.start_window)
+        self.assertEqual(self.marionette.current_chrome_window_handle, self.start_window)
+        self.marionette.close_chrome_window()
+        self.assertNotIn(self.start_window, self.marionette.chrome_window_handles)
+        self.assertEqual(self.marionette.current_chrome_window_handle, self.start_window)
+        self.assertEqual(len(self.marionette.chrome_window_handles), len(self.start_windows))
 
-        self.marionette.find_element(By.LINK_TEXT, "Open new window").click()
-        count = 0
-        while True:
-            window_handles = self.marionette.chrome_window_handles
-            window_handles.remove(current)
-            if len(window_handles) > 0:
-                break
-            elif count > 10:
-                self.fail("There were no windows that appeared when we clicked earlier")
-            else:
-                time.sleep(1)
-            count += 1
+    def test_should_load_and_close_a_window(self):
+        def open_window_with_link():
+            test_html = self.marionette.absolute_url("test_windows.html")
+            with self.marionette.using_context("content"):
+                self.marionette.navigate(test_html)
+                self.marionette.find_element(By.LINK_TEXT, "Open new window").click()
 
-        self.marionette.switch_to_window(window_handles[0])
-        self.assertEqual(self.marionette.title, "We Arrive Here")
+        new_window = self.open_window(trigger=open_window_with_link)
+        self.marionette.switch_to_window(new_window)
+        self.assertEqual(self.marionette.current_chrome_window_handle, new_window)
+        self.assertEqual(len(self.marionette.chrome_window_handles), 2)
 
-        handle = self.marionette.current_chrome_window_handle
-
-        self.assertEqual(self.marionette.current_chrome_window_handle, handle)
-        self.assertEqual(2, len(self.marionette.chrome_window_handles))
+        with self.marionette.using_context('content'):
+            self.assertEqual(self.marionette.title, "We Arrive Here")
 
         # Let's close and check
         self.marionette.close_chrome_window()
-        self.marionette.switch_to_window(current)
-        self.assertEqual(1, len(self.marionette.chrome_window_handles))
-
-    def testShouldCauseAWindowToLoadAndCheckItIsOpenThenCloseIt(self):
-        test_html = self.marionette.absolute_url("test_windows.html")
-        self.marionette.navigate(test_html)
-        current = self.marionette.current_chrome_window_handle
-
-        self.marionette.find_element(By.LINK_TEXT,"Open new window").click()
-        all_handles = self.marionette.chrome_window_handles
-        self.assertEqual(2, len(all_handles))
-        self.marionette.switch_to_window([x for x in all_handles if x != current][0])
-
-        # Let's close and check
-        self.marionette.close_chrome_window()
-        self.marionette.switch_to_window(current)
-        self.assertEqual(1, len(self.marionette.chrome_window_handles))
-
-    def tearDown(self):
-        #ensure that we close the window, regardless of pass/failure
-        self.close_new_window()
-        MarionetteTestCase.tearDown(self)
-
+        self.marionette.switch_to_window(self.start_window)
+        self.assertEqual(len(self.marionette.chrome_window_handles), 1)
--- a/testing/marionette/harness/marionette/tests/unit/test_window_switching.py
+++ b/testing/marionette/harness/marionette/tests/unit/test_window_switching.py
@@ -1,49 +1,44 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this file,
 # You can obtain one at http://mozilla.org/MPL/2.0/.
 
-from marionette import MarionetteTestCase
-
+from marionette import MarionetteTestCase, WindowManagerMixin
 from marionette_driver.by import By
 from marionette_driver.errors import NoSuchElementException
 from marionette_driver.wait import Wait
 
 
-class TestWindowSwitching(MarionetteTestCase):
+class TestWindowSwitching(WindowManagerMixin, MarionetteTestCase):
+
     def testJSWindowCreationAndSwitching(self):
         test_html = self.marionette.absolute_url("test_windows.html")
         self.marionette.navigate(test_html)
 
-        self.current_window = self.marionette.current_chrome_window_handle
-        link = self.marionette.find_element(By.LINK_TEXT, "Open new window")
-        link.click()
+        def open_window_with_link():
+            link = self.marionette.find_element(By.LINK_TEXT, "Open new window")
+            link.click()
 
-        windows = self.marionette.chrome_window_handles
-        windows.remove(self.current_window)
-        self.marionette.switch_to_window(windows[0])
+        new_window = self.open_window(trigger=open_window_with_link)
+        self.marionette.switch_to_window(new_window)
 
         title = self.marionette.execute_script("return document.title")
         results_page = self.marionette.absolute_url("resultPage.html")
         self.assertEqual(self.marionette.get_url(), results_page)
         self.assertEqual(title, "We Arrive Here")
 
-        #ensure navigate works in our current window
+        # ensure navigate works in our current window
         other_page = self.marionette.absolute_url("test.html")
         self.marionette.navigate(other_page)
-        other_window = self.marionette.current_chrome_window_handle
 
-        #try to access its dom
-        #since Bug 720714 stops us from checking DOMContentLoaded, we wait a bit
+        # try to access its dom
+        # since Bug 720714 stops us from checking DOMContentLoaded, we wait a bit
         Wait(self.marionette, timeout=30, ignored_exceptions=NoSuchElementException).until(
             lambda m: m.find_element(By.ID, 'mozLink'))
 
-        self.assertEqual(other_window, self.marionette.current_chrome_window_handle)
-        self.marionette.switch_to_window(self.current_window)
-        self.assertEqual(self.current_window, self.marionette.current_chrome_window_handle)
+        self.assertEqual(new_window, self.marionette.current_chrome_window_handle)
+        self.marionette.switch_to_window(self.start_window)
+        self.assertEqual(self.start_window, self.marionette.current_chrome_window_handle)
 
     def tearDown(self):
-        window_handles = self.marionette.chrome_window_handles
-        window_handles.remove(self.current_window)
-        for handle in window_handles:
-            self.marionette.switch_to_window(handle)
-            self.marionette.close_chrome_window()
+        self.close_all_windows()
+        super(TestWindowSwitching, self).tearDown()
--- a/testing/marionette/harness/marionette/www/windowHandles.html
+++ b/testing/marionette/harness/marionette/www/windowHandles.html
@@ -4,11 +4,11 @@
 
 <!DOCTYPE html>
 <html>
 <head>
 <title>Marionette New Tab Link</title>
 </head>
 <body>
   <a href="about:blank" id="new-tab" target="_blank">Click me!</a>
-  <a href="about:blank" id="new-window" onClick='javascript:window.open("about:blank", null, "location=1,toolbar=1");'>Click me!</a>
+  <a href="about:blank" id="new-window" onClick='javascript:window.open("windowHandles.html", null, "location=1,toolbar=1");'>Click me!</a>
 </body>
 </html>