take-visual.py

# 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 io
import os
import csv
import hashlib
import socket
import base64
import sys
import threading
import time
import traceback
import urlparse
import uuid
import logging

logging.basicConfig(format='%(asctime)s %(message)s',level=logging.INFO)

from PIL import Image

from selenium import webdriver
from selenium.common import exceptions
from selenium.webdriver.remote.remote_connection import RemoteConnection

wpt_root = os.path.join(os.path.dirname(__file__), "web-platform-tests")
sys.path[0:0] = [wpt_root]

from tools.manifest import manifest

extra_timeout = 5

size = (800, 600)

with open(os.path.join(wpt_root, "tools/wptrunner/wptrunner/executors/reftest-wait_webdriver.js")) as fp:
    wait = fp.read()

def resize_png(data, (width, height)):
    f = io.BytesIO(data)
    
    im = Image.open(f)
    old_width, old_height = im.size
    if old_width == width and old_height == height:
        return data
        logging.info("image is right size, rewriting anyway")
    else:
        logging.info("image was (%d, %d), resizing" % (old_width, old_height))

    mode = im.mode
    if len(mode) == 1:  # L, 1
        new_background = (255)
    if len(mode) == 3:  # RGB
        new_background = (255, 255, 255)
    if len(mode) == 4:  # RGBA, CMYK
        new_background = (255, 255, 255, 255)

    newImage = Image.new(mode, (width, height), new_background)
    newImage.paste(im, (0, 0, old_width, old_height))

    f2 = io.BytesIO()
    newImage.save(f2, "png")
    return f2.getvalue()
    
class Run(object):
    def __init__(self, func, webdriver, url, timeout):
        self.func = func
        self.result = None
        self.webdriver = webdriver
        self.url = url
        self.timeout = timeout
        self.result_flag = threading.Event()

    def run(self):
        timeout = self.timeout

        try:
            self.webdriver.set_script_timeout((timeout + extra_timeout) * 1000)
        except exceptions.ErrorInResponseException:
            logging.error("Lost WebDriver connection")
            return False, "Lost"

        executor = threading.Thread(target=self._run)
        executor.start()

        flag = self.result_flag.wait(timeout + 2 * extra_timeout)
        if self.result is None:
            assert not flag
            self.result = False, ("EXTERNAL-TIMEOUT", None)

        return self.result

    def _run(self):
        try:
            self.result = True, self.func(self.webdriver, self.url, self.timeout)
        except exceptions.TimeoutException:
            self.result = False, ("EXTERNAL-TIMEOUT", None)
        except (socket.timeout, exceptions.ErrorInResponseException):
            self.result = False, ("CRASH", None)
        except Exception as e:
            tb = traceback.format_exc()
            self.result = False, ("ERROR", tb)
        finally:
            self.result_flag.set()


class Executor(object):
    def __init__(self, webdriver):
        self.webdriver = webdriver

    def do_test(self, test):
        logging.info("run")
        return self.screenshot(test)

    def screenshot(self, test):
        return Run(self._screenshot,
                   self.webdriver,
                   test,
                   10).run()

    def _screenshot(self, webdriver, url, timeout):
        logging.info("shot")
        logging.info(url)

        webdriver.set_window_size(*size)
        
        webdriver.get(url)

        logging.info("shot1")
        rv = webdriver.execute_async_script(wait)
        logging.info("shot2")
        screenshot = webdriver.get_screenshot_as_base64()
        if screenshot.startswith("data:image/png;base64,"):
            screenshot = screenshot.split(",", 1)[1]
        screenshot = base64.b64decode(screenshot)
        logging.info("shot3")
        return screenshot

def main(browser):
    fp = open("out/%s.csv" % browser, "wb")
    fieldnames = ['url', 'sha256']
    writer = csv.DictWriter(fp, fieldnames=fieldnames)
    writer.writeheader()
    
    man = manifest.load("/", os.path.join(wpt_root, "MANIFEST.json"))

    ok = False
    for i, (_, _, path_tests) in enumerate(man.itertypes("visual")):
        if (i % 100) == 0:
            logging.info("restarting browser after 100 test paths")
            ok = False
        for test in path_tests:
            if not ok:
                try:
                    wd.quit()
                except:
                    pass
                if browser == "firefox":
                    profile = webdriver.FirefoxProfile()
                    # profile.set_preference('font.name.serif.x-western', 'Ahem')
                    # profile.set_preference('font.name.serif.x-unicode', 'Ahem')
                    # profile.set_preference('font.name.sans-serif.x-western', 'Ahem')
                    # profile.set_preference('font.name.sans-serif.x-unicode', 'Ahem')
                    # profile.set_preference('font.name.monospace.x-western', 'Ahem')
                    # profile.set_preference('font.name.monospace.x-unicode', 'Ahem')
                    # profile.set_preference('font.name.cursive.x-western', 'Ahem')
                    # profile.set_preference('font.name.cursive.x-unicode', 'Ahem')
                    # profile.set_preference('font.size.fixed.x-unicode', '16')
                    # profile.set_preference('font.size.fixed.x-western', '16')
                    # profile.set_preference('font.size.variable.x-unicode', '16')
                    # profile.set_preference('font.size.variable.x-western', '16')
                    wd = webdriver.Firefox(profile, firefox_binary="firefox-nightly")
                elif browser == "chrome":
                    options = webdriver.chrome.options.Options()
                    prefs = {
                    #     "webkit.webprefs.fonts.fixed.Zyyy": "Ahem",
                    #     "webkit.webprefs.fonts.sansserif.Zyyy": "Ahem",
                    #     "webkit.webprefs.fonts.serif.Zyyy": "Ahem",
                    #     "webkit.webprefs.fonts.standard.Zyyy": "Ahem",
                    #     "webkit.webprefs.default_font_size": "16",
                    #     "webkit.webprefs.default_fixed_font_size": "16"
                    }
                    options.add_experimental_option("prefs", prefs)
                    options.binary_location = '/usr/bin/google-chrome-unstable'
                    wd = webdriver.Chrome(chrome_options=options)
                elif browser == "safari":
                    wd = webdriver.Safari(executable_path="/Applications/Safari Technology Preview.app/Contents/MacOS/safaridriver")
                elif browser == "edge":
                    wd = webdriver.Edge()
                else:
                    raise Exception("foo")
                ex = Executor(wd)
                
            url = "http://web-platform.test:8000" + test.url
            logging.info("running %s" % url)
            ok, result = ex.do_test(url)
            if ok:
                # result = resize_png(result, size)
                logging.info("ok")
                m = hashlib.sha256()
                m.update(result)
                h = m.hexdigest()
                path = "out/screenshot/%s.png" % h
                if not os.path.exists(path):
                    with open(path, "wb") as fp2:
                        fp2.write(result)
                row = {}
                row["url"] = url
                row['sha256'] = h
                writer.writerow(row)
            else:
                logging.info("not ok, %s" % repr(result))
    try:
        wd.quit()
    except:
        pass

            
if __name__ == "__main__":
    main(sys.argv[1])