chore: screenshots task code cleanup (#5416)

Attempt to make the screenshot task more readable and easier to understand. The screenshot task includes a lot of magic (due to the Firebase functions JWT solution) and we should try our best to make it as readable as possible.

Author: devversion

tools/gulp/tasks/screenshots.ts

@@ -1,15 +1,26 @@
 import {task} from 'gulp';
 import {readdirSync, statSync, existsSync, mkdirp, readFileSync, writeFileSync} from 'fs-extra';
+import {openScreenshotsBucket, connectFirebaseScreenshots} from '../util/firebase';
+import {isTravisMasterBuild} from '../util/travis-ci';
+
 import * as path from 'path';
 import * as firebaseAdmin from 'firebase-admin';
-import * as firebase from 'firebase';
-import {
-  openScreenshotsBucket,
-  connectFirebaseScreenshots} from '../util/firebase';
-import {isTravisMasterBuild} from '../util/travis-ci';
 
+// Firebase provides TypeScript definitions that are only accessible from specific namespaces.
+// This means that those types are really long and it's nearly impossible to write a function that
+// doesn't exceed the maximum columns. Import the types from the namespace so they are shorter.
+import Database = firebaseAdmin.database.Database;
+import DataSnapshot = firebaseAdmin.database.DataSnapshot;
+
+// This import lacks of type definitions.
 const imageDiff = require('image-diff');
 
+/** Travis secure token that will be used by the Screenshot functions to verify the identity. */
+const travisSecureToken = getSecureToken();
+
+/** Git SHA of the current Pull Request being checked by Travis. */
+const pullRequestSha = process.env['TRAVIS_PULL_REQUEST_SHA'];
+
 const SCREENSHOT_DIR = './screenshots';
 const LOCAL_GOLDENS = path.join(SCREENSHOT_DIR, `golds`);
 const LOCAL_DIFFS = path.join(SCREENSHOT_DIR, `diff`);
@@ -29,157 +40,169 @@ task('screenshots', () => {
 
   if (isTravisMasterBuild()) {
     // Only update goldens for master build
-    return uploadScreenshots();
+    return uploadGoldenScreenshots();
   } else if (prNumber) {
     const firebaseApp = connectFirebaseScreenshots();
     const database = firebaseApp.database();
 
-    return updateTravis(database, prNumber)
-      .then(() => getScreenshotFiles(database))
-      .then(() => downloadAllGoldsAndCompare(database, prNumber))
-      .then((results: boolean) => updateResult(database, prNumber, results))
+    return uploadTravisJobInfo(database, prNumber)
+      .then(() => downloadGoldScreenshotFiles(database))
+      .then(() => compareScreenshotFiles(database, prNumber))
+      .then(passedAll => setPullRequestResult(database, prNumber, passedAll))
       .then(() => uploadScreenshotsData(database, 'diff', prNumber))
       .then(() => uploadScreenshotsData(database, 'test', prNumber))
       .catch((err: any) => console.error(err))
       .then(() => firebaseApp.delete());
   }
 });
 
-function updateFileResult(database: firebase.database.Database, prNumber: string,
-                          filenameKey: string, result: boolean) {
-  return getPullRequestRef(database, prNumber).child('results').child(filenameKey).set(result);
+/** Sets the screenshot diff result for a given file of a Pull Request. */
+function setFileResult(database: Database, prNumber: string, fileName: string, result: boolean) {
+  return getPullRequestRef(database, prNumber).child('results').child(fileName).set(result);
 }
 
-function updateResult(database: firebase.database.Database, prNumber: string, result: boolean) {
-  return getPullRequestRef(database, prNumber).child('result')
-    .child(process.env['TRAVIS_PULL_REQUEST_SHA']).set(result).then(() => result);
+/** Sets the full diff result for the current Pull Request that runs inside of Travis. */
+function setPullRequestResult(database: Database, prNumber: string, result: boolean) {
+  return getPullRequestRef(database, prNumber).child('result').child(pullRequestSha).set(result);
 }
 
-function getPullRequestRef(database: firebase.database.Database | firebaseAdmin.database.Database,
-                           prNumber: string) {
-  let secureToken = getSecureToken();
-  return database.ref(FIREBASE_REPORT).child(prNumber).child(secureToken);
+/** Returns the Firebase Reference that contains all data related to the specified PR. */
+function getPullRequestRef(database: Database, prNumber: string) {
+  return database.ref(FIREBASE_REPORT).child(prNumber).child(travisSecureToken);
 }
 
-function updateTravis(database: firebase.database.Database,
-                      prNumber: string) {
+/** Uploads necessary Travis CI job variables that will be used in the Screenshot Panel. */
+function uploadTravisJobInfo(database: Database, prNumber: string) {
   return getPullRequestRef(database, prNumber).update({
     sha: process.env['TRAVIS_PULL_REQUEST_SHA'],
     travis: process.env['TRAVIS_JOB_ID'],
   });
 }
 
-/** Get a list of filenames from firebase database. */
-function getScreenshotFiles(database: firebase.database.Database) {
+/** Downloads all golden screenshot files and stores them in the local file system. */
+function downloadGoldScreenshotFiles(database: Database) {
+  // Create the directory that will contain all goldens if it's not present yet.
   mkdirp(LOCAL_GOLDENS);
-  mkdirp(LOCAL_DIFFS);
-
-  return database.ref(FIREBASE_DATA_GOLDENS).once('value')
-      .then((snapshot: firebase.database.DataSnapshot) => {
-    let counter = 0;
-    snapshot.forEach((childSnapshot: firebase.database.DataSnapshot) => {
-      let key = childSnapshot.key;
-      let binaryData = new Buffer(childSnapshot.val(), 'base64').toString('binary');
-      writeFileSync(`${LOCAL_GOLDENS}/${key}.screenshot.png`, binaryData, 'binary');
-      counter++;
-      return counter == snapshot.numChildren();
+
+  return database.ref(FIREBASE_DATA_GOLDENS).once('value').then(snapshot => {
+    snapshot.forEach((childSnapshot: DataSnapshot) => {
+      const screenshotName = childSnapshot.key;
+      const binaryData = new Buffer(childSnapshot.val(), 'base64').toString('binary');
+
+      writeFileSync(`${LOCAL_GOLDENS}/${screenshotName}.screenshot.png`, binaryData, 'binary');
     });
-  }).catch((error: any) => console.log(error));
+  });
 }
 
+/** Extracts the name of a given screenshot file by removing the file extension. */
 function extractScreenshotName(fileName: string) {
   return path.basename(fileName, '.screenshot.png');
 }
 
-function getLocalScreenshotFiles(dir: string): string[] {
-  return readdirSync(dir)
-    .filter((fileName: string) => !statSync(path.join(dir, fileName)).isDirectory())
+/** Gets a list of files inside of a directory that end with `.screenshot.png`. */
+function getLocalScreenshotFiles(directory: string): string[] {
+  return readdirSync(directory)
+    .filter((fileName: string) => !statSync(path.join(directory, fileName)).isDirectory())
     .filter((fileName: string) => fileName.endsWith('.screenshot.png'));
 }
 
 /**
- * Get processed secure token. The jwt token has 3 parts: header, payload, signature and has format
- * {jwtHeader}.{jwtPayload}.{jwtSignature}
- * The three parts is connected by '.', while '.' is not a valid path in firebase database.
- * Replace all '.' to '/' to make the path valid
- * Output is {jwtHeader}/{jwtPayload}/{jwtSignature}.
- * This secure token is used to validate the write access is from our TravisCI under our repo.
- * All data is written to /$path/$secureToken/$data and after validated the
- * secure token, the data is moved to /$path/$data in database.
+ * Upload screenshots to a Firebase Database path that will then upload the file to a Google
+ * Cloud Storage bucket if the Auth token is valid.
+ * @param database Firebase database instance.
+ * @param prNumber The key used in firebase. Here it is the PR number.
+ * @param mode Upload mode. This can be either 'test' or 'diff'.
+ *  - If the images are the test results, mode should be 'test'.
+ *  - If the images are the diff images generated, mode should be 'diff'.
  */
-function getSecureToken() {
-  return process.env['FIREBASE_ACCESS_TOKEN'].replace(/[.]/g, '/');
-}
-
-/**
- * Upload screenshots to google cloud storage.
- * @param prNumber - The key used in firebase. Here it is the PR number.
- * @param mode - Can be 'test' or 'diff' .
- *   If the images are the test results, mode should be 'test'.
- *   If the images are the diff images generated, mode should be 'diff'.
- */
-function uploadScreenshotsData(database: firebase.database.Database,
-                               mode: 'test' | 'diff', prNumber: string) {
-  let localDir = mode == 'diff' ? path.join(SCREENSHOT_DIR, 'diff') : SCREENSHOT_DIR;
-  let promises = getLocalScreenshotFiles(localDir).map((file: string) => {
-    let fileName = path.join(localDir, file);
-    let filenameKey = extractScreenshotName(fileName);
-    let secureToken = getSecureToken();
-    let data = readFileSync(fileName);
-    return database.ref(FIREBASE_IMAGE).child(prNumber)
-      .child(secureToken).child(mode).child(filenameKey).set(data);
-  });
-  return Promise.all(promises);
+function uploadScreenshotsData(database: Database, mode: 'test' | 'diff', prNumber: string) {
+  const localDir = mode == 'diff' ? path.join(SCREENSHOT_DIR, 'diff') : SCREENSHOT_DIR;
+
+  return Promise.all(getLocalScreenshotFiles(localDir).map(file => {
+    const filePath = path.join(localDir, file);
+    const fileName = extractScreenshotName(filePath);
+    const binaryContent = readFileSync(filePath);
+
+    // Upload the Buffer of the screenshot image to a Firebase Database reference that will
+    // then upload the screenshot file to a Google Cloud Storage bucket if the JWT token is valid.
+    return database.ref(FIREBASE_IMAGE)
+      .child(prNumber).child(travisSecureToken).child(mode).child(fileName)
+      .set(binaryContent);
+  }));
 }
 
+/** Concurrently compares every golden screenshot with the newly taken screenshots. */
+function compareScreenshotFiles(database: Database, prNumber: string) {
+  const fileNames = getLocalScreenshotFiles(LOCAL_GOLDENS);
+  const compares = fileNames.map(fileName => compareScreenshotFile(fileName, database, prNumber));
 
-/** Download golds screenshots. */
-function downloadAllGoldsAndCompare(database: firebase.database.Database, prNumber: string) {
+  // Wait for all compares to finish and then return a Promise that resolves with a boolean that
+  // shows whether the tests passed or not.
+  return Promise.all(compares).then((results: boolean[]) => results.every(Boolean));
+}
 
-  let filenames = getLocalScreenshotFiles(LOCAL_GOLDENS);
+/** Compare the specified screenshot file with the golden file from Firebase. */
+function compareScreenshotFile(fileName: string, database: Database, prNumber: string) {
+  const goldScreenshotPath = path.join(LOCAL_GOLDENS, fileName);
+  const localScreenshotPath = path.join(SCREENSHOT_DIR, fileName);
+  const diffScreenshotPath = path.join(LOCAL_DIFFS, fileName);
 
-  return Promise.all(filenames.map((filename: string) => {
-    return diffScreenshot(filename, database, prNumber);
-  })).then((results: boolean[]) => results.every((value: boolean) => value == true));
-}
+  const screenshotName = extractScreenshotName(fileName);
 
-function diffScreenshot(filename: string, database: firebase.database.Database,
-                        prNumber: string) {
-  // TODO(tinayuangao): Run the downloads and diffs in parallel.
-  filename = path.basename(filename);
-  let goldUrl = path.join(LOCAL_GOLDENS, filename);
-  let pullRequestUrl = path.join(SCREENSHOT_DIR, filename);
-  let diffUrl = path.join(LOCAL_DIFFS, filename);
-  let filenameKey = extractScreenshotName(filename);
-
-  if (existsSync(goldUrl) && existsSync(pullRequestUrl)) {
-    return new Promise((resolve: any, reject: any) => {
-      imageDiff({
-        actualImage: pullRequestUrl,
-        expectedImage: goldUrl,
-        diffImage: diffUrl,
-      }, (err: any, imagesAreSame: boolean) => {
-        if (err) {
-          console.log(err);
-          imagesAreSame = false;
-          reject(err);
-        }
-        resolve(imagesAreSame);
-        return updateFileResult(database, prNumber, filenameKey, imagesAreSame);
+  if (existsSync(goldScreenshotPath) && existsSync(localScreenshotPath)) {
+    return compareImage(localScreenshotPath, goldScreenshotPath, diffScreenshotPath)
+      .then(result => {
+        // Set the screenshot diff result in Firebase and afterwards pass the result boolean
+        // to the Promise chain again.
+        return setFileResult(database, prNumber, screenshotName, result).then(() => result);
       });
-    });
   } else {
-    return updateFileResult(database, prNumber, filenameKey, false).then(() => false);
+    return setFileResult(database, prNumber, screenshotName, false).then(() => false);
   }
 }
 
-/** Upload screenshots to google cloud storage. */
-function uploadScreenshots() {
-  let bucket = openScreenshotsBucket();
-  let promises = getLocalScreenshotFiles(SCREENSHOT_DIR).map((file: string) => {
-    let fileName = path.join(SCREENSHOT_DIR, file);
-    let destination = `${FIREBASE_STORAGE_GOLDENS}/${file}`;
-    return bucket.upload(fileName, { destination: destination });
+/** Uploads golden screenshots to the Google Cloud Storage bucket for the screenshots. */
+function uploadGoldenScreenshots() {
+  const bucket = openScreenshotsBucket();
+
+  return Promise.all(getLocalScreenshotFiles(SCREENSHOT_DIR).map(fileName => {
+    const filePath = path.join(SCREENSHOT_DIR, fileName);
+    const storageDestination = `${FIREBASE_STORAGE_GOLDENS}/${filePath}`;
+
+    return bucket.upload(fileName, { destination: storageDestination });
+  }));
+}
+
+/**
+ * Compares two images using the Node package image-diff. A difference screenshot will be created.
+ * The returned promise will resolve with a boolean that will be true if the images are equal.
+ */
+function compareImage(actualPath: string, goldenPath: string, diffPath: string): Promise<boolean> {
+  return new Promise(resolve => {
+    imageDiff({
+      actualImage: actualPath,
+      expectedImage: goldenPath,
+      diffImage: diffPath,
+    }, (err: any, imagesAreEqual: boolean) => {
+      if (err) {
+        throw err;
+      }
+
+      resolve(imagesAreEqual);
+    });
   });
-  return Promise.all(promises);
+}
+
+/**
+ * Get processed secure token. The jwt token has 3 parts: header, payload, signature and has format
+ * {jwtHeader}.{jwtPayload}.{jwtSignature}
+ * The three parts is connected by '.', while '.' is not a valid path in firebase database.
+ * Replace all '.' to '/' to make the path valid
+ * Output is {jwtHeader}/{jwtPayload}/{jwtSignature}.
+ * This secure token is used to validate the write access is from our TravisCI under our repo.
+ * All data is written to /$path/$secureToken/$data and after validated the
+ * secure token, the data is moved to /$path/$data in database.
+ */
+function getSecureToken() {
+  return (process.env['FIREBASE_ACCESS_TOKEN'] || '').replace(/[.]/g, '/');
 }