snapshot-manager.js

'use strict';

const crypto = require('crypto');
const fs = require('fs');
const path = require('path');
const zlib = require('zlib');

const concordance = require('concordance');
const indentString = require('indent-string');
const md5Hex = require('md5-hex');
const convertSourceMap = require('convert-source-map');
const slash = require('slash');
const writeFileAtomic = require('write-file-atomic');
const mem = require('mem');

const concordanceOptions = require('./concordance-options').snapshotManager;

// Increment if encoding layout or Concordance serialization versions change. Previous AVA versions will not be able to
// decode buffers generated by a newer version, so changing this value will require a major version bump of AVA itself.
// The version is encoded as an unsigned 16 bit integer.
const VERSION = 2;

const VERSION_HEADER = Buffer.alloc(2);
VERSION_HEADER.writeUInt16LE(VERSION);

// The decoder matches on the trailing newline byte (0x0A).
const READABLE_PREFIX = Buffer.from(`AVA Snapshot v${VERSION}\n`, 'ascii');
const REPORT_SEPARATOR = Buffer.from('\n\n', 'ascii');
const REPORT_TRAILING_NEWLINE = Buffer.from('\n', 'ascii');

const MD5_HASH_LENGTH = 16;

class SnapshotError extends Error {
	constructor(message, snapPath) {
		super(message);
		this.name = 'SnapshotError';
		this.snapPath = snapPath;
	}
}
exports.SnapshotError = SnapshotError;

class ChecksumError extends SnapshotError {
	constructor(snapPath) {
		super('Checksum mismatch', snapPath);
		this.name = 'ChecksumError';
	}
}
exports.ChecksumError = ChecksumError;

class VersionMismatchError extends SnapshotError {
	constructor(snapPath, version) {
		super('Unexpected snapshot version', snapPath);
		this.name = 'VersionMismatchError';
		this.snapVersion = version;
		this.expectedVersion = VERSION;
	}
}
exports.VersionMismatchError = VersionMismatchError;

const LEGACY_SNAPSHOT_HEADER = Buffer.from('// Jest Snapshot v1');
function isLegacySnapshot(buffer) {
	return LEGACY_SNAPSHOT_HEADER.equals(buffer.slice(0, LEGACY_SNAPSHOT_HEADER.byteLength));
}

class LegacyError extends SnapshotError {
	constructor(snapPath) {
		super('Legacy snapshot file', snapPath);
		this.name = 'LegacyError';
	}
}
exports.LegacyError = LegacyError;

function tryRead(file) {
	try {
		return fs.readFileSync(file);
	} catch (error) {
		if (error.code === 'ENOENT') {
			return null;
		}

		throw error;
	}
}

function withoutLineEndings(buffer) {
	let checkPosition = buffer.byteLength - 1;
	while (buffer[checkPosition] === 0x0A || buffer[checkPosition] === 0x0D) {
		checkPosition--;
	}

	return buffer.slice(0, checkPosition + 1);
}

function formatEntry(label, descriptor) {
	if (label) {
		label = `> ${label}\n\n`;
	}

	const codeBlock = indentString(concordance.formatDescriptor(descriptor, concordanceOptions), 4);
	return Buffer.from(label + codeBlock, 'utf8');
}

function combineEntries(entries) {
	const buffers = [];
	let byteLength = 0;

	const sortedKeys = [...entries.keys()].sort((keyA, keyB) => {
		const [a, b] = [entries.get(keyA), entries.get(keyB)];
		const taskDifference = a.taskIndex - b.taskIndex;

		if (taskDifference !== 0) {
			return taskDifference;
		}

		const [assocA, assocB] = [a.associatedTaskIndex, b.associatedTaskIndex];
		if (assocA !== undefined && assocB !== undefined) {
			const assocDifference = assocA - assocB;

			if (assocDifference !== 0) {
				return assocDifference;
			}
		}

		return a.snapIndex - b.snapIndex;
	});

	for (const key of sortedKeys) {
		const keyBuffer = Buffer.from(`\n\n## ${key}\n\n`, 'utf8');
		buffers.push(keyBuffer);
		byteLength += keyBuffer.byteLength;

		const formattedEntries = entries.get(key).buffers;
		const last = formattedEntries[formattedEntries.length - 1];
		for (const entry of formattedEntries) {
			buffers.push(entry);
			byteLength += entry.byteLength;

			if (entry !== last) {
				buffers.push(REPORT_SEPARATOR);
				byteLength += REPORT_SEPARATOR.byteLength;
			}
		}
	}

	return {buffers, byteLength};
}

function generateReport(relFile, snapFile, entries) {
	const combined = combineEntries(entries);
	const {buffers} = combined;
	let {byteLength} = combined;

	const header = Buffer.from(`# Snapshot report for \`${slash(relFile)}\`

The actual snapshot is saved in \`${snapFile}\`.

Generated by [AVA](https://avajs.dev).`, 'utf8');
	buffers.unshift(header);
	byteLength += header.byteLength;

	buffers.push(REPORT_TRAILING_NEWLINE);
	byteLength += REPORT_TRAILING_NEWLINE.byteLength;
	return Buffer.concat(buffers, byteLength);
}

function appendReportEntries(existingReport, entries) {
	const combined = combineEntries(entries);
	const {buffers} = combined;
	let {byteLength} = combined;

	const prepend = withoutLineEndings(existingReport);
	buffers.unshift(prepend);
	byteLength += prepend.byteLength;

	buffers.push(REPORT_TRAILING_NEWLINE);
	byteLength += REPORT_TRAILING_NEWLINE.byteLength;
	return Buffer.concat(buffers, byteLength);
}

function encodeSnapshots(buffersByHash) {
	const buffers = [];
	let byteOffset = 0;

	// Entry start and end pointers are relative to the header length. This means
	// it's possible to append new entries to an existing snapshot file, without
	// having to rewrite pointers for existing entries.
	const headerLength = Buffer.alloc(4);
	buffers.push(headerLength);
	byteOffset += 4;

	// Allows 65535 hashes (tests or identified snapshots) per file.
	const numberHashes = Buffer.alloc(2);
	numberHashes.writeUInt16LE(buffersByHash.size);
	buffers.push(numberHashes);
	byteOffset += 2;

	const entries = [];
	// Maps can't have duplicate keys, so all items in [...buffersByHash.keys()]
	// are unique, so sortedHashes should be deterministic.
	const sortedHashes = [...buffersByHash.keys()].sort();
	const sortedBuffersByHash = [...sortedHashes.map(hash => [hash, buffersByHash.get(hash)])];
	for (const [hash, snapshotBuffers] of sortedBuffersByHash) {
		buffers.push(Buffer.from(hash, 'hex'));
		byteOffset += MD5_HASH_LENGTH;

		// Allows 65535 snapshots per hash.
		const numberSnapshots = Buffer.alloc(2);
		numberSnapshots.writeUInt16LE(snapshotBuffers.length, 0);
		buffers.push(numberSnapshots);
		byteOffset += 2;

		for (const value of snapshotBuffers) {
			// Each pointer is 32 bits, restricting the total, uncompressed buffer to
			// 4 GiB.
			const start = Buffer.alloc(4);
			const end = Buffer.alloc(4);
			entries.push({start, end, value});

			buffers.push(start, end);
			byteOffset += 8;
		}
	}

	headerLength.writeUInt32LE(byteOffset, 0);

	let bodyOffset = 0;
	for (const entry of entries) {
		const start = bodyOffset;
		const end = bodyOffset + entry.value.byteLength;
		entry.start.writeUInt32LE(start, 0);
		entry.end.writeUInt32LE(end, 0);
		buffers.push(entry.value);
		bodyOffset = end;
	}

	byteOffset += bodyOffset;

	const compressed = zlib.gzipSync(Buffer.concat(buffers, byteOffset));
	compressed[9] = 0x03; // Override the GZip header containing the OS to always be Linux
	const md5sum = crypto.createHash('md5').update(compressed).digest();
	return Buffer.concat([
		READABLE_PREFIX,
		VERSION_HEADER,
		md5sum,
		compressed
	], READABLE_PREFIX.byteLength + VERSION_HEADER.byteLength + MD5_HASH_LENGTH + compressed.byteLength);
}

function decodeSnapshots(buffer, snapPath) {
	if (isLegacySnapshot(buffer)) {
		throw new LegacyError(snapPath);
	}

	// The version starts after the readable prefix, which is ended by a newline
	// byte (0x0A).
	const versionOffset = buffer.indexOf(0x0A) + 1;
	const version = buffer.readUInt16LE(versionOffset);
	if (version !== VERSION) {
		throw new VersionMismatchError(snapPath, version);
	}

	const md5sumOffset = versionOffset + 2;
	const compressedOffset = md5sumOffset + MD5_HASH_LENGTH;
	const compressed = buffer.slice(compressedOffset);

	const md5sum = crypto.createHash('md5').update(compressed).digest();
	const expectedSum = buffer.slice(md5sumOffset, compressedOffset);
	if (!md5sum.equals(expectedSum)) {
		throw new ChecksumError(snapPath);
	}

	const decompressed = zlib.gunzipSync(compressed);
	let byteOffset = 0;

	const headerLength = decompressed.readUInt32LE(byteOffset);
	byteOffset += 4;

	const snapshotsByHash = new Map();
	const numberHashes = decompressed.readUInt16LE(byteOffset);
	byteOffset += 2;

	for (let count = 0; count < numberHashes; count++) {
		const hash = decompressed.toString('hex', byteOffset, byteOffset + MD5_HASH_LENGTH);
		byteOffset += MD5_HASH_LENGTH;

		const numberSnapshots = decompressed.readUInt16LE(byteOffset);
		byteOffset += 2;

		const snapshotsBuffers = new Array(numberSnapshots);
		for (let index = 0; index < numberSnapshots; index++) {
			const start = decompressed.readUInt32LE(byteOffset) + headerLength;
			byteOffset += 4;
			const end = decompressed.readUInt32LE(byteOffset) + headerLength;
			byteOffset += 4;
			snapshotsBuffers[index] = decompressed.slice(start, end);
		}

		// Allow for new entries to be appended to an existing header, which could
		// lead to the same hash being present multiple times.
		if (snapshotsByHash.has(hash)) {
			snapshotsByHash.set(hash, snapshotsByHash.get(hash).concat(snapshotsBuffers));
		} else {
			snapshotsByHash.set(hash, snapshotsBuffers);
		}
	}

	return snapshotsByHash;
}

class Manager {
	constructor(options) {
		this.appendOnly = options.appendOnly;
		this.dir = options.dir;
		this.recordNewSnapshots = options.recordNewSnapshots;
		this.relFile = options.relFile;
		this.reportFile = options.reportFile;
		this.snapFile = options.snapFile;
		this.snapPath = options.snapPath;
		this.snapshotsByHash = options.snapshotsByHash;

		this.hasChanges = false;
		this.reportEntries = new Map();
	}

	compare(options) {
		const hash = md5Hex(options.belongsTo);
		const entries = this.snapshotsByHash.get(hash) || [];
		const snapshotBuffer = entries[options.index];

		if (!snapshotBuffer) {
			if (!this.recordNewSnapshots) {
				return {pass: false};
			}

			if (options.deferRecording) {
				const record = this.deferRecord(hash, options);
				return {pass: true, record};
			}

			this.record(hash, options);
			return {pass: true};
		}

		const actual = concordance.deserialize(snapshotBuffer, concordanceOptions);
		const expected = concordance.describe(options.expected, concordanceOptions);
		const pass = concordance.compareDescriptors(actual, expected);

		return {actual, expected, pass};
	}

	deferRecord(hash, options) {
		const descriptor = concordance.describe(options.expected, concordanceOptions);
		const snapshot = concordance.serialize(descriptor);
		const entry = formatEntry(options.label, descriptor);
		const {taskIndex, snapIndex, associatedTaskIndex} = options;

		return () => { // Must be called in order!
			this.hasChanges = true;

			let snapshots = this.snapshotsByHash.get(hash);
			if (!snapshots) {
				snapshots = [];
				this.snapshotsByHash.set(hash, snapshots);
			}

			if (options.index > snapshots.length) {
				throw new RangeError(`Cannot record snapshot ${options.index} for ${JSON.stringify(options.belongsTo)}, exceeds expected index of ${snapshots.length}`);
			}

			if (options.index < snapshots.length) {
				throw new RangeError(`Cannot record snapshot ${options.index} for ${JSON.stringify(options.belongsTo)}, already exists`);
			}

			snapshots.push(snapshot);

			if (this.reportEntries.has(options.belongsTo)) {
				this.reportEntries.get(options.belongsTo).buffers.push(entry);
			} else {
				this.reportEntries.set(options.belongsTo, {buffers: [entry], taskIndex, snapIndex, associatedTaskIndex});
			}
		};
	}

	record(hash, options) {
		const record = this.deferRecord(hash, options);
		record();
	}

	save() {
		if (!this.hasChanges) {
			return null;
		}

		const {snapPath} = this;
		const buffer = encodeSnapshots(this.snapshotsByHash);

		const reportPath = path.join(this.dir, this.reportFile);
		const existingReport = this.appendOnly ? tryRead(reportPath) : null;
		const reportBuffer = existingReport ?
			appendReportEntries(existingReport, this.reportEntries) :
			generateReport(this.relFile, this.snapFile, this.reportEntries);

		fs.mkdirSync(this.dir, {recursive: true});

		const paths = [snapPath, reportPath];
		const tmpfileCreated = tmpfile => paths.push(tmpfile);
		writeFileAtomic.sync(snapPath, buffer, {tmpfileCreated});
		writeFileAtomic.sync(reportPath, reportBuffer, {tmpfileCreated});
		return paths;
	}
}

const resolveSourceFile = mem(file => {
	const testDir = path.dirname(file);
	const buffer = tryRead(file);
	if (!buffer) {
		return file; // Assume the file is stubbed in our test suite.
	}

	const source = buffer.toString();
	const converter = convertSourceMap.fromSource(source) || convertSourceMap.fromMapFileSource(source, testDir);
	if (converter) {
		const map = converter.toObject();
		const firstSource = `${map.sourceRoot || ''}${map.sources[0]}`;
		return path.resolve(testDir, firstSource);
	}

	return file;
});

const determineSnapshotDir = mem(({file, fixedLocation, projectDir}) => {
	const testDir = path.dirname(resolveSourceFile(file));
	if (fixedLocation) {
		const relativeTestLocation = path.relative(projectDir, testDir);
		return path.join(fixedLocation, relativeTestLocation);
	}

	const parts = new Set(path.relative(projectDir, testDir).split(path.sep));
	if (parts.has('__tests__')) {
		return path.join(testDir, '__snapshots__');
	}

	if (parts.has('test') || parts.has('tests')) { // Accept tests, even though it's not in the default test patterns
		return path.join(testDir, 'snapshots');
	}

	return testDir;
}, {cacheKey: ([{file}]) => file});

exports.determineSnapshotDir = determineSnapshotDir;

function load({file, fixedLocation, projectDir, recordNewSnapshots, updating}) {
	const dir = determineSnapshotDir({file, fixedLocation, projectDir});
	const relFile = path.relative(projectDir, resolveSourceFile(file));
	const name = path.basename(relFile);
	const reportFile = `${name}.md`;
	const snapFile = `${name}.snap`;
	const snapPath = path.join(dir, snapFile);

	let appendOnly = !updating;
	let snapshotsByHash;

	if (!updating) {
		const buffer = tryRead(snapPath);
		if (buffer) {
			snapshotsByHash = decodeSnapshots(buffer, snapPath);
		} else {
			appendOnly = false;
		}
	}

	return new Manager({
		appendOnly,
		dir,
		recordNewSnapshots,
		relFile,
		reportFile,
		snapFile,
		snapPath,
		snapshotsByHash: snapshotsByHash || new Map()
	});
}

exports.load = load;