Upload saucelabs screenshots to depicted for comparison

saucelabs_to_depicted.py

import sys

from optparse import make_option

from StringIO import StringIO

from PIL import Image
import requests

from django.conf import settings
from django.core.management.base import BaseCommand


SAUCELABS_USERNAME = settings.SAUCELABS_USERNAME
SAUCELABS_ACCESS_KEY = settings.SAUCELABS_ACCESS_KEY
SAUCELABS_API_ROOT = 'https://saucelabs.com/rest/v1/%s/jobs' % SAUCELABS_USERNAME

# Depicted API server. Currently using the test server ran by Depicted author.
# Could be your own instance.
DEPICTED_API_ID = settings.DEPICTED_API_ID
DEPICTED_API_SECRET = settings.DEPICTED_API_SECRET
DEPICTED_BUILD_ID = settings.DEPICTED_BUILD_ID
DEPICTED_URL = "https://dpxdt-test.appspot.com"
DEPICTED_API_ROOT = '%s/api/' % DEPICTED_URL


class Command(BaseCommand):
    """
    The purpose of this script is to upload saucelabs screenshots from our
    ApplicationSanityTests to depicted for comparison.

    You provide a travis build number, and it will grab sauce jobs, find the
    correct job for that travis build number, and upload the assets to
    depicted.
    """
    help = 'Upload photos from sauce runs to depicted.'

    option_list = BaseCommand.option_list + (
        make_option('-b', '--build',
            action='store',
            dest='build',
            help="Travis build number."),
    )

    def handle(self, *args, **options):

        travis_build = options.get('build')

        if not travis_build:
            raise AttributeError("Travis build number required.")

        saucelabs_job_id = self.get_sauce_job(travis_build)

        # Get job assets
        r = requests.get('%s/%s/assets' % (SAUCELABS_API_ROOT, saucelabs_job_id),
            auth=(SAUCELABS_USERNAME, SAUCELABS_ACCESS_KEY))

        screenshots = r.json()['screenshots']

        # Upload each screenshot to depicted.
        print 'Uploading %s screenshots.' % len(screenshots)
        uploads = []
        for screenshot in screenshots:

            sha = self.upload_screenshot(screenshot, saucelabs_job_id)
            uploads.append({
                'sha': sha,
                'name': screenshot
            })
        print '\nScreenshots uploaded.'

        # Create depicted release
        release_name = 'Sprintly Travis build #%s' % travis_build
        release_number = self.create_depicted_release(release_name, travis_build)

        # Report runs. Each run is a test comparing a screenshot.
        print 'Report depicted runs.'
        for upload in uploads:
            self.report_depicted_run(upload, release_name, release_number)

        print '\nTests reported. Done.'
        print "Find the result at: %s/build?id=%s" % (
            DEPICTED_URL, DEPICTED_BUILD_ID)

    def get_sauce_job(self, travis_build):
        """Return saucelabs job from travis build number."""

        # Almost 30 sauce jobs per travis run.
        print 'Fetching saucelabs jobs.'
        r = requests.get('%s?limit=150' % SAUCELABS_API_ROOT, auth=(SAUCELABS_USERNAME, SAUCELABS_ACCESS_KEY))

        # Look for ApplicationSanityTests test class with this build number.
        # We use ApplicationSanityTests because they do not change often,
        # any change would most likely be a regression.
        job_id = None
        for job in r.json():
            r = requests.get('%s/%s' % (SAUCELABS_API_ROOT, job['id']),
                auth=(SAUCELABS_USERNAME, SAUCELABS_ACCESS_KEY))

            job_payload = r.json()

            # Is this the right job?
            if job_payload['build'] == travis_build and "ApplicationSanityTests" in job_payload['name']:
                job_id = job['id']
                print '\nFound saucelabs job.'
                break
            sys.stdout.write('.')
            sys.stdout.flush()

        if not job_id:
            # Handle case when you're using a Travis job that is too old.
            raise Exception('Could not find matching job. Maybe it was too long ago?')

        return job_id

    def upload_screenshot(self, screenshot, saucelabs_job_id):
        """
        Get screenshot contents from saucelabs API and upload it to depicted.
        Return upload sha from depicted.
        """

        # Get image contents.
        r = requests.get('%s/%s/assets/%s' % (SAUCELABS_API_ROOT, saucelabs_job_id, screenshot),
            auth=(SAUCELABS_USERNAME, SAUCELABS_ACCESS_KEY), stream=True)

        image = Image.open(StringIO(r.raw.read()))

        w, h = image.size
        image = image.crop((0, 22, w, h))

        # Debug what the image looks like here.
        # image.show()

        # Create empty file to save to
        fp = StringIO()
        image.save(fp, format='PNG')
        fp.seek(0)

        # Post file to depicted
        files = {'file': (screenshot, fp, 'image/png')}
        r = requests.post('%supload' % DEPICTED_API_ROOT, files=files,
            data={'build_id': DEPICTED_BUILD_ID},
            auth=(DEPICTED_API_ID, DEPICTED_API_SECRET))

        sys.stdout.write('.')
        sys.stdout.flush()

        return r.json()['sha1sum']

    def create_depicted_release(self, release_name, travis_build):
        """Create a depicted release for this build. These should mirror Travis job IDs"""

        print 'Create depicted release.'
        release = {
            'build_id': DEPICTED_BUILD_ID,
            'release_name': release_name,
            'url': 'https://sprint.ly'
        }
        r = requests.post('%screate_release' % DEPICTED_API_ROOT,
            data=release,
            auth=(DEPICTED_API_ID, DEPICTED_API_SECRET))
        release_number = r.json()['release_number'] 
        print 'Depicted release #%s created.' % release_number
        return release_number

    def report_depicted_run(self, upload, release_name, release_number):
        """
        Create depicted tests. Try to compare them to last good tests.
        
        This Depicted endpoint name is confusing. Basically the report_run
        endpoint is used to run image comparison. First we find the last
        'good' run and then trigger a new run telling it to compare to the
        last good image.
        """

        # Find the last good run for this test
        run = {
            'build_id': DEPICTED_BUILD_ID,
            'run_name': upload['name'],
        }
        r = requests.post('%sfind_run' % DEPICTED_API_ROOT,
            data=run,
            auth=(DEPICTED_API_ID, DEPICTED_API_SECRET), stream=True)

        try:
            last_good_image_sha = r.json()['image']
        except KeyError:
            # Baseline run?
            last_good_image_sha = None

        # Create a new run, comparing to the last good image
        run = {
            'build_id': DEPICTED_BUILD_ID,
            'release_name': release_name,
            'release_number': release_number,
            'url': 'https://sprint.ly',
            'run_name': upload['name'],
            'image': upload['sha'],
        }

        if last_good_image_sha:
            run['ref_image'] = last_good_image_sha

        r = requests.post('%sreport_run' % DEPICTED_API_ROOT,
            data=run,
            auth=(DEPICTED_API_ID, DEPICTED_API_SECRET), stream=True)

        if r.json()['success']:
            sys.stdout.write('.')
            sys.stdout.flush()
        else:
            print r.json()