Use conditional typing for defining the right set of options

src/main.ts

@@ -20,7 +20,11 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 // SOFTWARE.
 
-import { CloudFunction as CloudFunctionV1 } from 'firebase-functions';
+import {
+  CloudFunction as CloudFunctionV1,
+  HttpsFunction,
+  Runnable,
+} from 'firebase-functions';
 
 import {
   CloudFunction as CloudFunctionV2,
@@ -31,6 +35,11 @@ import { wrapV1, WrappedFunction, WrappedScheduledFunction } from './v1';
 
 import { wrapV2, WrappedV2Function } from './v2';
 
+type HttpsFunctionOrCloudFunctionV1<T, U> = U extends HttpsFunction &
+  Runnable<T>
+  ? HttpsFunction & Runnable<T>
+  : CloudFunctionV1<T>;
+
 // Re-exporting V1 (to reduce breakage)
 export {
   ContextOptions,
@@ -45,20 +54,24 @@ export {
 // V2 Exports
 export { WrappedV2Function } from './v2';
 
+export function wrap<T>(
+  cloudFunction: HttpsFunction & Runnable<T>
+): WrappedFunction<T, HttpsFunction & Runnable<T>>;
 export function wrap<T>(
   cloudFunction: CloudFunctionV1<T>
 ): WrappedScheduledFunction | WrappedFunction<T>;
 export function wrap<T extends CloudEvent<unknown>>(
   cloudFunction: CloudFunctionV2<T>
 ): WrappedV2Function<T>;
-
 export function wrap<T, V extends CloudEvent<unknown>>(
   cloudFunction: CloudFunctionV1<T> | CloudFunctionV2<V>
 ): WrappedScheduledFunction | WrappedFunction<T> | WrappedV2Function<V> {
   if (isV2CloudFunction<V>(cloudFunction)) {
     return wrapV2<V>(cloudFunction as CloudFunctionV2<V>);
   }
-  return wrapV1<T>(cloudFunction as CloudFunctionV1<T>);
+  return wrapV1<T>(
+    cloudFunction as HttpsFunctionOrCloudFunctionV1<T, typeof cloudFunction>
+  );
 }
 
 /**

src/v1.ts

@@ -30,6 +30,8 @@ import {
   config,
   database,
   firestore,
+  HttpsFunction,
+  Runnable,
 } from 'firebase-functions';
 
 /** Fields of the event context that can be overridden/customized. */
@@ -88,14 +90,16 @@ export type CallableContextOptions = {
 };
 
 /* Fields for both Event and Callable contexts, checked at runtime */
-export type ContextOptions = EventContextOptions | CallableContextOptions;
+export type ContextOptions<T = void> = T extends HttpsFunction & Runnable<T>
+  ? CallableContextOptions
+  : EventContextOptions;
 
 /** A function that can be called with test data and optional override values for the event context.
  * It will subsequently invoke the cloud function it wraps with the provided test data and a generated event context.
  */
-export type WrappedFunction<T> = (
+export type WrappedFunction<T, U = void> = (
   data: T,
-  options?: ContextOptions
+  options?: ContextOptions<U>
 ) => any | Promise<any>;
 
 /** A scheduled function that can be called with optional override values for the event context.
@@ -106,9 +110,12 @@ export type WrappedScheduledFunction = (
 ) => any | Promise<any>;
 
 /** Takes a cloud function to be tested, and returns a WrappedFunction which can be called in test code. */
+export function wrapV1<T>(
+  cloudFunction: HttpsFunction & Runnable<T>
+): WrappedFunction<T, HttpsFunction & Runnable<T>>;
 export function wrapV1<T>(
   cloudFunction: CloudFunction<T>
-): WrappedScheduledFunction | WrappedFunction<T> {
+): WrappedScheduledFunction | WrappedFunction<T, CloudFunction<T>> {
   if (!has(cloudFunction, '__trigger')) {
     throw new Error(
       'Wrap can only be called on functions written with the firebase-functions SDK.'
@@ -150,26 +157,26 @@ export function wrapV1<T>(
   const isCallableFunction =
     get(cloudFunction, '__trigger.labels.deployment-callable') === 'true';
 
-  let wrapped: WrappedFunction<T> = (data: T, options: ContextOptions) => {
+  let wrapped: WrappedFunction<T, typeof cloudFunction> = (data, options) => {
     // Although in Typescript we require `options` some of our JS samples do not pass it.
-    options = options || {};
+    const _options = { ...options };
     let context;
 
     if (isCallableFunction) {
       _checkOptionValidity(
         ['app', 'auth', 'instanceIdToken', 'rawRequest'],
-        options
+        _options
       );
-      let callableContextOptions = options as CallableContextOptions;
+      let callableContextOptions = _options as CallableContextOptions;
       context = {
         ...callableContextOptions,
       };
     } else {
       _checkOptionValidity(
         ['eventId', 'timestamp', 'params', 'auth', 'authType', 'resource'],
-        options
+        _options
       );
-      const defaultContext = _makeDefaultContext(cloudFunction, options, data);
+      const defaultContext = _makeDefaultContext(cloudFunction, _options, data);
 
       if (
         has(defaultContext, 'eventType') &&
@@ -179,7 +186,7 @@ export function wrapV1<T>(
         defaultContext.authType = 'UNAUTHENTICATED';
         defaultContext.auth = null;
       }
-      context = merge({}, defaultContext, options);
+      context = merge({}, defaultContext, _options);
     }
 
     return cloudFunction.run(data, context);
@@ -226,9 +233,14 @@ function _checkOptionValidity(
   });
 }
 
+function _makeDefaultContext<T>(
+  cloudFunction: HttpsFunction & Runnable<T>,
+  options: CallableContextOptions,
+  triggerData?: T
+);
 function _makeDefaultContext<T>(
   cloudFunction: CloudFunction<T>,
-  options: ContextOptions,
+  options: EventContextOptions,
   triggerData?: T
 ): EventContext {
   let eventContextOptions = options as EventContextOptions;