Stub Lambda Functions in your CDK Stack

Readme.md

Stubbing Lambda Functions in your CDK Stack

Given you have a stack with one or more Lambda functions (e.g. as part of a Step Functions state machine), it can be pretty useful to stub long running parts with a known response.

This makes use of cdk Aspects, which allows modifying all or a filtered subsset of resources for a given scope (Stack, Construct).

In addition this leverages raw overrides to remove the original code of the Lambda function.

Please note, that the stub has to be in Python or NodeJS, since inlining code is only supported by those runtimes.

lambda-stub.ts

import { IAspect, IConstruct, Construct } from '@aws-cdk/core';
import { InlineCode, Function, CfnFunction, Runtime} from '@aws-cdk/aws-lambda';

class StubManager {
  public static isStubbable(construct: any): construct is Function {
    return construct.constructor.name === "Function"
  }
}

export interface LambdaStubProps {
  code: InlineCode;
  runTime: Runtime;
  handler: string;
}

abstract class LambdaStubBase implements IAspect {
  public readonly lambdaName: string;
  protected readonly props: LambdaStubProps;

  constructor(lambaName: string, props: LambdaStubProps) {
    this.lambdaName = lambaName;
    this.props = props;
  }

  public visit(construct: IConstruct): void {        
    if (StubManager.isStubbable(construct)) {
      const fn = construct as Function
      if (fn.node.id === this.lambdaName) {
        this.applyStub(fn)
      }
    }
  }

  protected abstract applyStub(resource: Function): void;
}

export class LambdaStub extends LambdaStubBase {
  public readonly lambdaName: string
  public readonly code: InlineCode

  public static add(scope: Construct, lambdaName: string, props: LambdaStubProps) {
    scope.node.applyAspect(new LambdaStub(lambdaName, props));
  }

  constructor(lambdaName: string, props: LambdaStubProps) {
    super(lambdaName, props);
    this.code = props.code  
    this.lambdaName = lambdaName
  }

  protected applyStub(resource: Function): void {
    const cnfFunction = resource.node.defaultChild as CfnFunction
    cnfFunction.addPropertyOverride("Runtime", this.props.runTime.name)
    cnfFunction.addPropertyOverride("Handler", this.props.handler)
    cnfFunction.addPropertyOverride("Code", {
      ZipFile: this.code.bind(cnfFunction).inlineCode,
      S3Bucket: undefined,
      S3Key: undefined
    })
    cnfFunction.addDeletionOverride("Metadata.aws:asset:property")
    cnfFunction.addDeletionOverride("Metadata.aws:asset:path")
  }
}

step-functions-stack.ts

import cdk = require('@aws-cdk/core');
import lambda = require('@aws-cdk/aws-lambda');
const path = require('path')

export class StepFunctionsStack extends cdk.Stack {  
  constructor(scope: cdk.Construct, id: string, props?: cdk.StackProps) {
    super(scope, id, props);

    new lambda.Function(this, 'Some Complex Workflow', {
      code: lambda.Code.asset(path.resolve(__dirname, 'lambda', 'complex-stuff')),
      handler: 'index.handler',
      timeout: cdk.Duration.minutes(15),
      runtime: lambda.Runtime.RUBY_2_5
    });
    
    // more stuff ... 
  }
}

stubbed-stack.ts

#!/usr/bin/env node
import 'source-map-support/register';
import cdk = require('@aws-cdk/core');
import { StepFunctionsStack } from './step-functions-stack';
import { LambdaStub } from './lambda'
import * as lambda from '@aws-cdk/aws-lambda'

const app = new cdk.App();
const stack  = new StepFunctionsStack(app, 'StepTest');

LambdaStub.add(stack, 'Some Complex Workflow', {
  handler: 'index.handler',
  code: lambda.Code.fromInline('exports.handler = async function(event, ctx) { return "success" }'),
  runTime: lambda.Runtime.NODEJS_12_X
})