feat: implement :header: and screenshot_default_headers (#45)

.github/workflows/unit-test.yml

@@ -27,11 +27,28 @@ jobs:
       - name: Playwright install
         run: playwright install
       - name: Run unittest
-        run: pytest
+        id: pytest
+        shell: python
+        run: |
+          import os
+          import subprocess
+          import sys
+
+          result = subprocess.run(['pytest'], capture_output=True, text=True)
+          output = result.stdout
+          print(output)
+          if result.returncode > 0:
+            files = [line.strip() for line in output.splitlines() if "obtained" in line and os.path.exists(line.strip())]
+            if files:
+              print("Faulty files:", files)
+              with open(os.getenv("GITHUB_OUTPUT"), "a") as fd:
+                fd.write("imgs=" + "\n".join(files))
+
+          sys.exit(result.returncode)
       - name: Upload failed image artifact
         if: failure()
         uses: actions/upload-artifact@v4
         with:
           name: failed-images-${{ matrix.os }}-${{ matrix.python-version }}
-          # https://docs.pytest.org/en/stable/how-to/tmp_path.html#temporary-directory-location-and-retention
-          path: /tmp/pytest-of-*/pytest-*/*/*/*.obtained.png
+          path: |
+            ${{ steps.pytest.outputs.imgs }}

sphinxcontrib/screenshot.py

@@ -107,6 +107,7 @@ class ScreenshotDirective(SphinxDirective):
       'color-scheme': str,
       'full-page': directives.flag,
       'context': str,
+      'headers': directives.unchanged,
   }
   pool = ThreadPoolExecutor()
 
@@ -116,7 +117,8 @@ def take_screenshot(
       init_script: str, interactions: str, generate_pdf: bool,
       color_scheme: ColorScheme, full_page: bool,
       context_builder: typing.Optional[typing.Callable[[Browser, str, str],
-                                                       BrowserContext]]):
+                                                       BrowserContext]],
+      headers: dict):
     """Takes a screenshot with Playwright's Chromium browser.
 
     Args:
@@ -153,6 +155,7 @@ def take_screenshot(
       try:
         if init_script:
           page.add_init_script(init_script)
+        page.set_extra_http_headers(headers)
         page.goto(url)
         page.wait_for_load_state('networkidle')
 
@@ -201,6 +204,13 @@ def run(self) -> typing.List[nodes.Node]:
                  self.env.config.screenshot_default_full_page)
     context = self.options.get('context', '')
     interactions = '\n'.join(self.content)
+    headers = self.options.get('headers', '')
+
+    request_headers = {**self.env.config.screenshot_default_headers}
+    if headers:
+      for header in headers.strip().split("\n"):
+        name, value = header.split(" ", 1)
+        request_headers[name] = value
 
     if urlparse(url).scheme not in {'http', 'https'}:
       raise RuntimeError(
@@ -227,7 +237,7 @@ def run(self) -> typing.List[nodes.Node]:
       fut = self.pool.submit(ScreenshotDirective.take_screenshot, url, browser,
                              width, height, filepath, screenshot_init_script,
                              interactions, pdf, color_scheme, full_page,
-                             context_builder)
+                             context_builder, request_headers)
       fut.result()
 
     # Create image and figure nodes
@@ -317,6 +327,10 @@ def setup(app: Sphinx) -> Meta:
       'env',
       types=[dict[str, str]],
       description="A dict of paths to Playwright context build methods")
+  app.add_config_value(
+      'screenshot_default_headers', {},
+      'env',
+      description="The default headers to pass in requests")
   app.add_config_value(
       'screenshot_apps', {},
       'env',

tests/roots/test-headers/conf.py

@@ -0,0 +1,21 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     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.
+import pathlib
+import sys
+
+sys.path.insert(0, str(pathlib.Path(__file__).parent.resolve()))
+
+extensions = ['sphinxcontrib.screenshot']
+screenshot_apps = {"example": "example_headers_app:create_app"}
+screenshot_default_headers = {"Accept-Language": "fr-FR,fr"}

tests/roots/test-headers/example_headers_app.py

@@ -0,0 +1,33 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     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.
+def create_app(sphinx_app):
+
+  def hello_world_app(environ, start_response):
+    headers = [('Content-type', 'text/html; charset=utf-8')]
+    start_response('200 OK', headers)
+    style = (b"font-family: arial; " + b"font-size: 30px; " + b"padding: 0; " +
+             b"margin: 0; " + b"font-smooth: never; " + b"line-height: 1.15;")
+
+    accept_encoding = environ.get('HTTP_ACCEPT_ENCODING', 'unset').encode()
+    accept_language = environ.get('HTTP_ACCEPT_LANGUAGE', 'unset').encode()
+    authorization = environ.get('HTTP_AUTHORIZATION', 'unset').encode()
+
+    return [
+        b"<html>" + b'<body style="' + style + b'">' + b"Accept-Encoding: " +
+        accept_encoding + b"<br>" + b"Accept-Language: " + accept_language +
+        b"<br>" + b"Authorization: " + authorization + b"<br>"
+        b"</body>" + b"</html>"
+    ]
+
+  return hello_world_app

tests/roots/test-headers/index.rst

@@ -0,0 +1,6 @@
+.. screenshot:: |example|
+  :width: 800
+  :height: 600
+  :headers:
+    Authorization Bearer foobar
+    Accept-Encoding gzip, deflate, br

tests/test_headers.py

@@ -0,0 +1,33 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     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 io import StringIO
+
+import pytest
+from bs4 import BeautifulSoup
+from sphinx.testing.util import SphinxTestApp
+
+
+@pytest.mark.sphinx('html', testroot="headers")
+def test_headers(app: SphinxTestApp, status: StringIO, warning: StringIO,
+                 image_regression) -> None:
+  app.build()
+  out_html = app.outdir / "index.html"
+
+  soup = BeautifulSoup(out_html.read_text(), "html.parser")
+  imgs = soup.find_all('img')
+
+  img_path = app.outdir / imgs[0]['src']
+  with open(img_path, "rb") as fd:
+    # Tolerance of 0.5%
+    image_regression.check(fd.read(), diff_threshold=0.5)