1 files changed, 85 insertions, 8 deletions

diff --git a/test/lib/content-function.js b/test/lib/content-function.js
index b51f2847..21af0e5a 100644
--- a/test/lib/content-function.js
+++ b/test/lib/content-function.js
@@ -7,11 +7,15 @@ import urlSpec from '../../src/url-spec.js';
 import {getColors} from '../../src/util/colors.js';
 import {generateURLs} from '../../src/util/urls.js';
 
+import mock from './generic-mock.js';
+
 export function testContentFunctions(t, message, fn) {
   const urls = generateURLs(urlSpec);
 
   t.test(message, async t => {
-    const loadedContentDependencies = await quickLoadContentDependencies();
+    let loadedContentDependencies;
+
+    const mocks = [];
 
     const evaluate = ({
       from = 'localized.home',
@@ -19,9 +23,13 @@ export function testContentFunctions(t, message, fn) {
       extraDependencies = {},
       ...opts
     }) => {
+      if (!loadedContentDependencies) {
+        throw new Error(`Await .load() before performing tests`);
+      }
+
       const {to} = urls.from(from);
 
-      try {
+      return cleanCatchAggregate(() => {
         return quickEvaluate({
           ...opts,
           contentDependencies: {
@@ -37,19 +45,88 @@ export function testContentFunctions(t, message, fn) {
             ...extraDependencies,
           },
         });
-      } catch (error) {
-        if (error instanceof AggregateError) {
-          error = new Error(`AggregateError: ${error.message}\n${error.errors.map(err => `** ${err}`).join('\n')}`);
-        }
-        throw error;
+      });
+    };
+
+    evaluate.load = async (opts) => {
+      if (loadedContentDependencies) {
+        throw new Error(`Already loaded!`);
       }
+
+      loadedContentDependencies = await asyncCleanCatchAggregate(() =>
+        quickLoadContentDependencies(opts));
     };
 
     evaluate.snapshot = (opts, fn) => {
+      if (!loadedContentDependencies) {
+        throw new Error(`Await .load() before performing tests`);
+      }
+
       const result = (fn ? fn(evaluate(opts)) : evaluate(opts));
       t.matchSnapshot(result.toString(), 'output');
     };
 
-    return fn(t, evaluate);
+    evaluate.mock = (...opts) => {
+      const {value, close} = mock(...opts);
+      mocks.push({close});
+      return value;
+    };
+
+    await fn(t, evaluate);
+
+    if (mocks.length) {
+      cleanCatchAggregate(() => {
+        const errors = [];
+        for (const {close} of mocks) {
+          try {
+            close();
+          } catch (error) {
+            errors.push(error);
+          }
+        }
+        if (errors.length) {
+          throw new AggregateError(errors, `Errors closing mocks`);
+        }
+      });
+    }
   });
 }
+
+function cleanAggregate(error) {
+  if (error instanceof AggregateError) {
+    return new Error(`[AggregateError: ${error.message}\n${
+      error.errors
+        .map(cleanAggregate)
+        .map(err => ` * ${err.message.split('\n').map((l, i) => (i > 0 ? '   ' + l : l)).join('\n')}`)
+        .join('\n')}]`);
+  } else {
+    return error;
+  }
+}
+
+function printAggregate(error) {
+  if (error instanceof AggregateError) {
+    const {message} = cleanAggregate(error);
+    for (const line of message.split('\n')) {
+      console.error(line);
+    }
+  }
+}
+
+function cleanCatchAggregate(fn) {
+  try {
+    return fn();
+  } catch (error) {
+    printAggregate(error);
+    throw error;
+  }
+}
+
+async function asyncCleanCatchAggregate(fn) {
+  try {
+    return await fn();
+  } catch (error) {
+    printAggregate(error);
+    throw error;
+  }
+}