1import requests
2from mozlog.structured.formatters.base import BaseFormatter
3
4DEFAULT_API = "https://wpt.fyi/api/screenshots/hashes"
5
6
7class WptscreenshotFormatter(BaseFormatter):  # type: ignore
8    """Formatter that outputs screenshots in the format expected by wpt.fyi."""
9
10    def __init__(self, api=None):
11        self.api = api or DEFAULT_API
12        self.cache = set()
13
14    def suite_start(self, data):
15        # TODO(Hexcles): We might want to move the request into a different
16        # place, make it non-blocking, and handle errors better.
17        params = {}
18        run_info = data.get("run_info", {})
19        if "product" in run_info:
20            params["browser"] = run_info["product"]
21        if "browser_version" in run_info:
22            params["browser_version"] = run_info["browser_version"]
23        if "os" in run_info:
24            params["os"] = run_info["os"]
25        if "os_version" in run_info:
26            params["os_version"] = run_info["os_version"]
27        try:
28            r = requests.get(self.api, params=params)
29            r.raise_for_status()
30            self.cache = set(r.json())
31        except (requests.exceptions.RequestException, ValueError):
32            pass
33
34    def test_end(self, data):
35        if "reftest_screenshots" not in data.get("extra", {}):
36            return
37        output = ""
38        for item in data["extra"]["reftest_screenshots"]:
39            if type(item) != dict:
40                # Skip the relation string.
41                continue
42            checksum = "sha1:" + item["hash"]
43            if checksum in self.cache:
44                continue
45            self.cache.add(checksum)
46            output += "data:image/png;base64,{}\n".format(item["screenshot"])
47        return output if output else None
48