如何使用Jest模拟AWS Lambda中的lambda.invoke().promise()调用？

问题描述 #

我有一段待测试的单元代码：

import { Lambda } from 'aws-sdk'

export const handler = async (
  event: APIGatewayEvent,
  ctx: Context,
  cb: APIGatewayProxyCallback
): Promise<void> => {
  
  const lambda = new Lambda()
  await lambda.invoke(params).promise()
}

我希望在Jest中模拟其中的await lambda.invoke(params).promise()调用，让它直接resolve或“跳过”，但尝试使用jest.spyOn未成功，请问该如何实现？

解决方案 #

方法一：直接mock整个AWS SDK的Lambda类 #

通过jest.mock直接替换aws-sdk模块中的Lambda实现，让它返回预设的mock方法：

// 测试文件
import { handler } from './your-handler-file';
import { Lambda } from 'aws-sdk';

// 模拟aws-sdk的Lambda类
jest.mock('aws-sdk', () => {
  // 模拟invoke方法，链式调用需要返回自身
  const mockInvoke = jest.fn().mockReturnThis();
  // 模拟promise方法，让它直接resolve
  const mockPromise = jest.fn().mockResolvedValue({});

  return {
    Lambda: jest.fn(() => ({
      invoke: mockInvoke,
      promise: mockPromise
    }))
  };
});

describe('handler 单元测试', () => {
  it('应正确模拟lambda.invoke().promise()调用', async () => {
    await handler({}, {} as Context, jest.fn());

    // 验证调用逻辑
    expect(Lambda).toHaveBeenCalled();
    expect(Lambda.mock.instances[0].invoke).toHaveBeenCalled(); // 可添加预期的params参数
    expect(Lambda.mock.instances[0].promise).toHaveBeenCalled();
  });
});

方法二：使用jest.spyOn模拟原型方法 #

因为代码中是通过new Lambda()创建实例，所以需要Spy Lambda类原型上的方法，而非类本身：

import { handler } from './your-handler-file';
import { Lambda } from 'aws-sdk';

describe('handler 单元测试', () => {
  let invokeSpy: jest.SpyInstance;
  let promiseSpy: jest.SpyInstance;

  beforeEach(() => {
    // 模拟invoke方法，链式调用需返回this
    invokeSpy = jest.spyOn(Lambda.prototype, 'invoke').mockReturnThis();
    // 让promise方法直接resolve
    promiseSpy = jest.spyOn(Lambda.prototype, 'promise').mockResolvedValue({});
  });

  afterEach(() => {
    // 恢复原方法，避免影响其他测试
    invokeSpy.mockRestore();
    promiseSpy.mockRestore();
  });

  it('应正确跳过lambda调用逻辑', async () => {
    await handler({}, {} as Context, jest.fn());

    expect(invokeSpy).toHaveBeenCalled();
    expect(promiseSpy).toHaveBeenCalled();
  });
});

为什么之前的jest.spyOn失败？ #

大概率是因为你直接Spy了Lambda类而非它的原型对象，或者没有处理链式调用：invoke()需要返回一个带有promise()方法的对象（通过mockReturnThis()实现），否则链式调用会报错，导致mock失效。

内容的提问来源于stack exchange，提问作者Leo