diff --git a/types/index.d.ts b/types/index.d.ts
index 1c4fdea..2cc9ab8 100644
--- a/types/index.d.ts
+++ b/types/index.d.ts
@@ -17,10 +17,11 @@ interface Warmer {
    * @param event the event passed to the lambda
    * @param [config] the config options to change lambda warmer's default behavior.  All of
    * the settings are optional
+   * @param [context] the context passed to the lambda
    *
    * @returns a Promise that resolves to true if this is a warming invocation
    */
-  (event: any, config?: WarmerConfig): Promise<boolean>;
+  (event: any, config?: WarmerConfig, context?: any): Promise<boolean>;
   WarmerConfig: WarmerConfig;
 }
 
diff --git a/types/test.ts b/types/test.ts
index 7855ed1..ddd7788 100644
--- a/types/test.ts
+++ b/types/test.ts
@@ -2,6 +2,7 @@ import warmer from 'lambda-warmer';
 
 warmer({});
 warmer({}, {});
+warmer({}, {}, {});
 warmer({}, {
   flag: 'some-flag',
   concurrency: 10,
@@ -10,6 +11,8 @@ warmer({}, {
   correlationId: 'some-id',
   delay: 1000,
   target: 'some-target',
+}, {
+  invokedFunctionArn: 'some-arn'
 });
 
 // @ts-expect-error - params should be required