Fix for screenshots with function scoped browsers

pytest_splinter/plugin.py

@@ -336,6 +336,107 @@ def splinter_clean_cookies_urls():
     return []
 
 
+def _take_screenshot(
+        request,
+        browser_instance,
+        fixture_name,
+        splinter_screenshot_dir,
+        splinter_screenshot_getter_html,
+        splinter_screenshot_getter_png
+):
+    """Capture a screenshot as .png and .html.
+
+    Invoked from session and function browser fixtures.
+    """
+    slaveoutput = getattr(request.config, 'slaveoutput', None)
+    try:
+        names = junitxml.mangle_testnames(request.node.nodeid.split("::"))
+    except AttributeError:
+        # pytest>=2.9.0
+        names = junitxml.mangle_test_address(request.node.nodeid)
+
+    classname = '.'.join(names[:-1])
+    screenshot_dir = os.path.join(splinter_screenshot_dir, classname)
+    screenshot_file_name_format = '{0}.{{format}}'.format(
+        '{0}-{1}'.format(names[-1][:128 - len(fixture_name) - 5], fixture_name).replace(os.path.sep, '-')
+    )
+    screenshot_file_name = screenshot_file_name_format.format(format='png')
+    screenshot_html_file_name = screenshot_file_name_format.format(format='html')
+    if not slaveoutput:
+        if not os.path.exists(screenshot_dir):
+            os.makedirs(screenshot_dir)
+    else:
+        screenshot_dir = session_tmpdir.ensure('screenshots', dir=True).strpath
+    screenshot_png_path = os.path.join(screenshot_dir, screenshot_file_name)
+    screenshot_html_path = os.path.join(screenshot_dir, screenshot_html_file_name)
+    LOGGER.info('Saving screenshot to %s', screenshot_dir)
+    try:
+        splinter_screenshot_getter_html(browser_instance, screenshot_html_path)
+        splinter_screenshot_getter_png(browser_instance, screenshot_png_path)
+        if request.node.splinter_failure.longrepr:
+            reprtraceback = request.node.splinter_failure.longrepr.reprtraceback
+            reprtraceback.extraline = _screenshot_extraline(screenshot_png_path, screenshot_html_path)
+        if slaveoutput is not None:
+            with codecs.open(screenshot_html_path, encoding=splinter_screenshot_encoding) as html_fd:
+                with open(screenshot_png_path) as fd:
+                    slaveoutput.setdefault('screenshots', []).append({
+                        'class_name': classname,
+                        'files': [
+                            {
+                                'file_name': screenshot_file_name,
+                                'content': fd.read(),
+                            },
+                            {
+                                'file_name': screenshot_html_file_name,
+                                'content': html_fd.read(),
+                                'encoding': splinter_screenshot_encoding
+                            }]
+                    })
+    except Exception as e:  # NOQA
+        request.config.warn('SPL504', "Could not save screenshot: {0}".format(e))
+
+
+@pytest.yield_fixture(autouse=True)
+def _browser_screenshot_session(
+        request,
+        splinter_session_scoped_browser,
+        splinter_screenshot_dir,
+        splinter_make_screenshot_on_failure,
+        splinter_screenshot_getter_html,
+        splinter_screenshot_getter_png
+):
+    """Make browser screenshot on test failure."""
+    yield
+
+    # Screenshot for function scoped browsers is handled in browser_instance_getter
+    if not splinter_session_scoped_browser:
+        return
+
+    fixture_values = (
+        # pytest 3
+        getattr(request, '_fixture_values', {}) or
+        # pytest 2
+        getattr(request, '_funcargs', {})
+    )
+
+    for name, value in fixture_values.items():
+        should_take_screenshot = (
+            hasattr(value, '__splinter_browser__') and
+            splinter_make_screenshot_on_failure and
+            request.node.splinter_failure
+        )
+
+        if should_take_screenshot:
+            _take_screenshot(
+                request=request,
+                fixture_name=name,
+                browser_instance=value,
+                splinter_screenshot_dir=splinter_screenshot_dir,
+                splinter_screenshot_getter_html=splinter_screenshot_getter_html,
+                splinter_screenshot_getter_png=splinter_screenshot_getter_png,
+            )
+
+
 @pytest.fixture(scope='session')
 def browser_instance_getter(
         browser_patches,
@@ -358,6 +459,9 @@ def browser_instance_getter(
         splinter_window_size,
         splinter_browser_class,
         splinter_clean_cookies_urls,
+        splinter_screenshot_getter_html,
+        splinter_screenshot_getter_png,
+        splinter_screenshot_encoding,
         session_tmpdir,
         browser_pool,
 ):
@@ -398,6 +502,20 @@ def prepare_browser(request, parent):
                 request.addfinalizer(browser.quit)
         elif not browser:
             browser = browser_pool[browser_key] = get_browser(splinter_webdriver)
+
+        if request.scope == 'function':
+            def _take_screenshot_on_failure():
+                if splinter_make_screenshot_on_failure and request.node.splinter_failure:
+                    _take_screenshot(
+                        request=request,
+                        fixture_name=parent.__name__,
+                        browser_instance=browser,
+                        splinter_screenshot_dir=splinter_screenshot_dir,
+                        splinter_screenshot_getter_html=splinter_screenshot_getter_html,
+                        splinter_screenshot_getter_png=splinter_screenshot_getter_png,
+                    )
+            request.addfinalizer(_take_screenshot_on_failure)
+
         try:
             if splinter_webdriver not in browser.driver_name.lower():
                 raise IOError('webdriver does not match')
@@ -430,68 +548,6 @@ def prepare_browser(request, parent):
     return prepare_browser
 
 
-@pytest.yield_fixture(autouse=True)
-def browser_screenshot(
-        request, splinter_screenshot_dir, session_tmpdir, splinter_make_screenshot_on_failure,
-        splinter_screenshot_encoding, splinter_screenshot_getter_png, splinter_screenshot_getter_html):
-    """Make browser screenshot on test failure."""
-    yield
-    for name, value in (
-            # pytest 3
-            getattr(request, '_fixture_values', {}) or
-            # pytest 2
-            getattr(request, '_funcargs', {})).items():
-        if hasattr(value, '__splinter_browser__'):
-            browser = value
-            if splinter_make_screenshot_on_failure and request.node.splinter_failure:
-                slaveoutput = getattr(request.config, 'slaveoutput', None)
-                try:
-                    names = junitxml.mangle_testnames(request.node.nodeid.split("::"))
-                except AttributeError:
-                    # pytest>=2.9.0
-                    names = junitxml.mangle_test_address(request.node.nodeid)
-
-                classname = '.'.join(names[:-1])
-                screenshot_dir = os.path.join(splinter_screenshot_dir, classname)
-                screenshot_file_name_format = '{0}.{{format}}'.format(
-                    '{0}-{1}'.format(names[-1][:128 - len(name) - 5], name).replace(os.path.sep, '-')
-                )
-                screenshot_file_name = screenshot_file_name_format.format(format='png')
-                screenshot_html_file_name = screenshot_file_name_format.format(format='html')
-                if not slaveoutput:
-                    if not os.path.exists(screenshot_dir):
-                        os.makedirs(screenshot_dir)
-                else:
-                    screenshot_dir = session_tmpdir.ensure('screenshots', dir=True).strpath
-                screenshot_png_path = os.path.join(screenshot_dir, screenshot_file_name)
-                screenshot_html_path = os.path.join(screenshot_dir, screenshot_html_file_name)
-                LOGGER.info('Saving screenshot to %s', screenshot_dir)
-                try:
-                    splinter_screenshot_getter_html(browser, screenshot_html_path)
-                    splinter_screenshot_getter_png(browser, screenshot_png_path)
-                    if request.node.splinter_failure.longrepr:
-                        reprtraceback = request.node.splinter_failure.longrepr.reprtraceback
-                        reprtraceback.extraline = _screenshot_extraline(screenshot_png_path, screenshot_html_path)
-                    if slaveoutput is not None:
-                        with codecs.open(screenshot_html_path, encoding=splinter_screenshot_encoding) as html_fd:
-                            with open(screenshot_png_path) as fd:
-                                slaveoutput.setdefault('screenshots', []).append({
-                                    'class_name': classname,
-                                    'files': [
-                                        {
-                                            'file_name': screenshot_file_name,
-                                            'content': fd.read(),
-                                        },
-                                        {
-                                            'file_name': screenshot_html_file_name,
-                                            'content': html_fd.read(),
-                                            'encoding': splinter_screenshot_encoding
-                                        }]
-                                })
-                except Exception as e:  # NOQA
-                    request.config.warn('SPL504', "Could not save screenshot: {0}".format(e))
-
-
 @pytest.hookimpl(hookwrapper=True)
 def pytest_runtest_makereport(item, call):
     """Assign the report to the item for futher usage."""

tests/test_plugin.py

@@ -154,6 +154,14 @@ def test_executable():
     assert arg2['executable_path'] == '/tmp'
 
 
+def assert_valid_html_screenshot_content(content):
+    """Make sure content fetched from html screenshoting looks correct."""
+    assert content.startswith('<html xmlns="http://www.w3.org/1999/xhtml">')
+    assert '<div id="content">' in content
+    assert '<strong>text</strong>' in content
+    assert content.endswith('</html>')
+
+
 def test_browser_screenshot_normal(testdir, simple_page_content):
     """Test making screenshots on test failure.
 
@@ -172,8 +180,35 @@ def test_screenshot(simple_page, browser):
             assert False
     """.format(simple_page_content), "-vl", "-r w")
 
+    assert testdir.tmpdir.join('test_browser_screenshot_normal', 'test_screenshot-browser.png')
     content = testdir.tmpdir.join('test_browser_screenshot_normal', 'test_screenshot-browser.html').read()
-    assert content.replace('\n', '') == simple_page_content.replace('\n', '')
+    assert_valid_html_screenshot_content(content)
+
+
+def test_browser_screenshot_function_scoped_browser(testdir, simple_page_content):
+    """Test making screenshots on test failure.
+
+    Normal test run.
+    """
+    testdir.inline_runsource("""
+        import pytest
+
+        @pytest.fixture
+        def simple_page(httpserver, browser):
+            httpserver.serve_content(
+                '''{0}''', code=200, headers={{'Content-Type': 'text/html'}})
+            browser.visit(httpserver.url)
+
+        def test_screenshot(simple_page, browser):
+            assert False
+    """.format(simple_page_content), "-vl", "-r w", '--splinter-session-scoped-browser=false')
+
+    content = testdir.tmpdir.join(
+        'test_browser_screenshot_function_scoped_browser',
+        'test_screenshot-browser.html'
+    ).read()
+
+    assert_valid_html_screenshot_content(content)
     assert testdir.tmpdir.join('test_browser_screenshot_normal', 'test_screenshot-browser.png')
 
 
@@ -198,5 +233,5 @@ def test_screenshot(simple_page, browser, param):
 
     content = testdir.tmpdir.join(
         'test_browser_screenshot_escaped', 'test_screenshot[escaped-param]-browser.html').read()
-    assert content.replace('\n', '') == simple_page_content.replace('\n', '')
+    assert_valid_html_screenshot_content(content)
     assert testdir.tmpdir.join('test_browser_screenshot_escaped', 'test_screenshot[escaped-param]-browser.png')