UI Screenshot Comparison Lambda Function

screenshotCompare.js

'use strict';
console.log('Loading function');

var AWS = require('aws-sdk');
var fs = require('fs');
var https = require('https');
var spawn = require('child_process').spawn;
var util = require('util');

AWS.config.apiVersions = {
    s3: '2006-03-01'
};

exports.handler = function(event, context, done) {
    var mainEvent = event.Records[0].s3;
    console.log(mainEvent);

    var bucketName = mainEvent.bucket.name;
    var s3 = new AWS.S3({params: {Bucket: bucketName}});

    var uploadedKey = decodeURIComponent(mainEvent.object.key);
    var uploadedKeyParts = uploadedKey.split('/');
    var suffix = uploadedKeyParts.pop();
    var prNumber = uploadedKeyParts.pop();
    var masterCompare = 'screenshots/master/' + suffix;

    // Don't do anything if we're uploading a master screenshot or a diff
    if (uploadedKey.indexOf('/master/') !== -1 || uploadedKey.indexOf('/diff-') !== -1) {
        return;
    }

    var receivedNew, receivedMaster = null;

    var makeComment = function(diffKey) {
        var title = suffix.split('.')[0].replace(/-/g, ' ').replace('compare ', '');
        var imgPrefix = `![](https://s3.amazonaws.com/${bucketName}/`;
        var imgPostfix = `?version=${mainEvent.object.eTag})`;
        var headers = ['Original', 'This PR'];
        var headerSeparator = '---|---';
        var images = [
            `${imgPrefix}${masterCompare}${imgPostfix}`,
            `${imgPrefix}${uploadedKey}${imgPostfix}`,
        ];
        
        if (diffKey) {
            headers.push('Diff');
            headerSeparator += '|---';
            images.push(`${imgPrefix}${diffKey}${imgPostfix}`);
        }
        
        var commentParts = [
            `<h4>Please check the "${title}" UI element for intended changes</h4>`,
            headers.join(' | '),
            headerSeparator,
            images.join(' | ')
        ];


        if (!diffKey) {
            commentParts.push(
                '<p>Unable to create diff image because the size of the screenshots changed.</p>'
            );
        }
        var commentBody = commentParts.join('\n') + '\n';

        var githubApiComment = {
            hostname: 'api.github.com',
            path: `/repos/[Owner]/[Repo]/issues/${prNumber}/comments`,
            method: 'POST',
            auth: '[bot]:[token]',
            headers: {
                'user-agent': 'Lambda-UI-Comparisons'
            }
        };

        var commentRequest = https.request(githubApiComment, function(response) {
            console.log(`GITHUB COMMENT STATUS: ${response.statusCode}`);
            console.log(`GITHUB COMMENT MESSAGE: ${response.statusMessage}`);
            done(null, 'Success');
        });
        commentRequest.on('error', function(err) {
            console.log(err);
        });
        commentRequest.write(JSON.stringify({body: commentBody}));
        commentRequest.end();
    };

    var compareImages = function() {
        var diffPath = '/tmp/diff-' + suffix;

        var stdout = '';
        var stderr = '';

        var compareProc = spawn('compare', [
            '-metric', 'mse',
            receivedMaster, receivedNew,
            //'-highlight-color', 'yellow',
            diffPath
        ]);

        compareProc.stdout.on('data', function(data) { stdout+=data });
        compareProc.stderr.on('data', function(data) { stderr+=data });
        compareProc.on('close', function (code) {
            // Difference information goes to stderr
            if (stderr) {
                var regex = /\((\d+\.?[\d\-\+e]*)\)/m;
                var match = regex.exec(stderr);
                if (stderr.indexOf('widths or heights differ') !== -1) {
                    return makeComment(false);
                }
                if (!match) {
                    return console.log('Unable to parse compare output.', stdout, stderr);
                }

                var equality = parseFloat(match[1]);
                if (equality > 0) {
                    console.log(equality);
                    // Upload the diff img to s3
                    var diffKey = uploadedKey.replace('compare', 'diff');
                    var diffStream = fs.createReadStream(diffPath);
                    var uploadData = {
                        Key: diffKey,
                        Body: diffStream,
                        ContentType: 'image/png'
                    };
                    s3.putObject(uploadData, function(err, data) {
                        if (err) {
                            return console.log('s3 upload error', err);
                        }
                        makeComment(diffKey);
                    });
                }
            }
            else {
                return console.log('compare stdout', stdout);
            }
        });
    }

    // Get the uploaded file and the master comparison file
    // Then when both requests are completed, do the comparison
    var tmp = '/tmp/';
    var uploadedPath = tmp + suffix;
    var masterPath = tmp + 'master-' + suffix;

    var uploadStream = fs.createWriteStream(uploadedPath);
    var masterStream = fs.createWriteStream(masterPath);

    var logError = function(err) {
        console.log('s3 download error', err);
    };

    s3.getObject({Key: uploadedKey}).on('error', logError).createReadStream().pipe(uploadStream).on('close', function() {
        receivedNew = uploadedPath;

        if (receivedNew && receivedMaster) {
            // we have both images, we can continue processing
            compareImages();
        }
    });
    s3.getObject({Key: masterCompare}).on('error', logError).createReadStream().pipe(masterStream).on('close', function() {
        receivedMaster = masterPath;

        if (receivedNew && receivedMaster) {
            // we have both images, we can continue processing
            compareImages();
        }
    });
};