Add support for scheduled functions. (#51)

Author: rhodgkins

spec/index.spec.ts

@@ -72,6 +72,7 @@ import './app.spec';
 import './providers/https.spec';
 import './providers/firestore.spec';
 import './providers/database.spec';
+import './providers/scheduled.spec';
 // import './providers/analytics.spec';
 // import './providers/auth.spec';
 // import './providers/https.spec';

spec/providers/scheduled.spec.ts

@@ -0,0 +1,51 @@
+import * as sinon from 'sinon';
+import * as functions from 'firebase-functions';
+import fft = require('../../src/index');
+import { WrappedScheduledFunction } from '../../src/main';
+
+describe('providers/scheduled', () => {
+  const fakeFn = sinon.fake.resolves();
+  const scheduledFunc = functions.pubsub
+    .schedule('every 2 hours')
+    .onRun(fakeFn);
+
+  const emptyObjectMatcher = sinon.match(
+    (v) => sinon.match.object.test(v) && Object.keys(v).length === 0
+  );
+
+  afterEach(() => {
+    fakeFn.resetHistory();
+  });
+
+  it('should run the wrapped function with generated context', async () => {
+    const test = fft();
+    const fn: WrappedScheduledFunction = test.wrap(scheduledFunc);
+    await fn();
+    // Function should only be called with 1 argument
+    sinon.assert.calledOnce(fakeFn);
+    sinon.assert.calledWithExactly(
+      fakeFn,
+      sinon.match({
+        eventType: sinon.match.string,
+        timestamp: sinon.match.string,
+        params: emptyObjectMatcher,
+      })
+    );
+  });
+
+  it('should run the wrapped function with provided context', async () => {
+    const timestamp = new Date().toISOString();
+    const test = fft();
+    const fn: WrappedScheduledFunction = test.wrap(scheduledFunc);
+    await fn({ timestamp });
+    sinon.assert.calledOnce(fakeFn);
+    sinon.assert.calledWithExactly(
+      fakeFn,
+      sinon.match({
+        eventType: sinon.match.string,
+        timestamp,
+        params: emptyObjectMatcher,
+      })
+    );
+  });
+});

src/main.ts

@@ -73,14 +73,40 @@ export type WrappedFunction = (
   options?: ContextOptions
 ) => any | Promise<any>;
 
+/** A scheduled function that can be called with optional override values for the event context.
+ * It will subsequently invoke the cloud function it wraps with a generated event context.
+ */
+export type WrappedScheduledFunction = (
+  options?: ContextOptions
+) => any | Promise<any>;
+
 /** Takes a cloud function to be tested, and returns a WrappedFunction which can be called in test code. */
-export function wrap<T>(cloudFunction: CloudFunction<T>): WrappedFunction {
+export function wrap<T>(
+  cloudFunction: CloudFunction<T>
+): WrappedScheduledFunction | WrappedFunction {
   if (!has(cloudFunction, '__trigger')) {
     throw new Error(
       'Wrap can only be called on functions written with the firebase-functions SDK.'
     );
   }
 
+  if (get(cloudFunction, '__trigger.labels.deployment-scheduled') === 'true') {
+    const scheduledWrapped: WrappedScheduledFunction = (
+      options: ContextOptions
+    ) => {
+      // Although in Typescript we require `options` some of our JS samples do not pass it.
+      options = options || {};
+
+      _checkOptionValidity(['eventId', 'timestamp'], options);
+      const defaultContext = _makeDefaultContext(cloudFunction, options);
+      const context = merge({}, defaultContext, options);
+
+      // @ts-ignore
+      return cloudFunction.run(context);
+    };
+    return scheduledWrapped;
+  }
+
   if (
     has(cloudFunction, '__trigger.httpsTrigger') &&
     get(cloudFunction, '__trigger.labels.deployment-callable') !== 'true'
@@ -115,20 +141,7 @@ export function wrap<T>(cloudFunction: CloudFunction<T>): WrappedFunction {
         ['eventId', 'timestamp', 'params', 'auth', 'authType'],
         options
       );
-      let eventContextOptions = options as EventContextOptions;
-      const defaultContext: EventContext = {
-        eventId: _makeEventId(),
-        resource: cloudFunction.__trigger.eventTrigger && {
-          service: cloudFunction.__trigger.eventTrigger.service,
-          name: _makeResourceName(
-            cloudFunction.__trigger.eventTrigger.resource,
-            has(eventContextOptions, 'params') && eventContextOptions.params
-          ),
-        },
-        eventType: get(cloudFunction, '__trigger.eventTrigger.eventType'),
-        timestamp: new Date().toISOString(),
-        params: {},
-      };
+      const defaultContext = _makeDefaultContext(cloudFunction, options);
 
       if (
         has(defaultContext, 'eventType') &&
@@ -138,7 +151,7 @@ export function wrap<T>(cloudFunction: CloudFunction<T>): WrappedFunction {
         defaultContext.authType = 'UNAUTHENTICATED';
         defaultContext.auth = null;
       }
-      context = merge({}, defaultContext, eventContextOptions);
+      context = merge({}, defaultContext, options);
     }
 
     return cloudFunction.run(data, context);
@@ -185,6 +198,27 @@ function _checkOptionValidity(
   });
 }
 
+function _makeDefaultContext<T>(
+  cloudFunction: CloudFunction<T>,
+  options: ContextOptions
+): EventContext {
+  let eventContextOptions = options as EventContextOptions;
+  const defaultContext: EventContext = {
+    eventId: _makeEventId(),
+    resource: cloudFunction.__trigger.eventTrigger && {
+      service: cloudFunction.__trigger.eventTrigger.service,
+      name: _makeResourceName(
+        cloudFunction.__trigger.eventTrigger.resource,
+        has(eventContextOptions, 'params') && eventContextOptions.params
+      ),
+    },
+    eventType: get(cloudFunction, '__trigger.eventTrigger.eventType'),
+    timestamp: new Date().toISOString(),
+    params: {},
+  };
+  return defaultContext;
+}
+
 /** Make a Change object to be used as test data for Firestore and real time database onWrite and onUpdate functions. */
 export function makeChange<T>(before: T, after: T): Change<T> {
   return Change.fromObjects(before, after);