AWS Lambda: The structured approach to simple and testable code

Published on: Sun Aug 21 2022

Jerry Chang

Goals

In this module, we will use this starter template - aws-webhook-series-tracing-and-logging as a starting point.

In this article, we’ll answer questions such as:

✅ How to structure our code so that the code is testable ?
✅ What are the steps to simple and testable code ?
✅ How and what tests should we be writing ?

We will use the example of the AWS Lambda error handler to demonstrate the design process, implementation and testing.

Content

Introduction
Design the structure
Error handler utility function
Final integration
1. Integrate handleError into lambda function
2. Tweak the Lambda tests
Conclusion

Introduction

In Build a webhook microservice using AWS techincal series, we designed the infrastructure and wrote the functional code for our solution.

Before we started coding, we started with the scaffold of the individual “units” then progressed the integrated solutions by combining these lego blocks.

Let’s do a quick review of our design.

Here is the scaffold we started with:

Ingestion function scaffold with comments

Illustration of the scaffold of the lambda function with comments

Individual units:

Illustration of the “units” in the Lambda function

So, that means we will be creating a module that will take care of our error handling in our lambda function.

Let’s dig into the details of what that would look like.

Designing the structure

Ideally, we should have a centralized error handler. It should be self-contained inside its own module.

So, when we use the error handler, we won’t need to know about all its details, we just know that if we give it an error object, it should return with the error response.

Illustration of the possible outcomes from the Lambda function

This has few benefits:

Testability - It is easier to unit test, rather than having the error handling logic all over the place, it is self-contained in a module
Readability - It it much easier to reason about the outcomes of the lambda function
Extensibility - It is easier to extend the error handling logic when it is self contained (ie adding custom logging and metrics)

These are just a few benefits.

The biggest benefit is the testability (or making it easier to test).

💡 Helpful: What I found helps is that if you focus on makng your modules or functions easy to test, in the process, you usually make the code more modular.
This is because when your module or function is doing too many things at once, it is more difficult to effectively test it. So, decomposing it to smaller units is required.

What would that look like in our code ?

Here is a scaffold of what that may look like:

// src/index.ts

export const handler = async(
  event: APIGatewayProxyEvent
): Promise<APIGatewayProxyResult> => {
  captureRequestContext(event);
  let messageId: string = '';
  try {
    // 1. Verify Signature
    verifySignature(event);
    // 2. Add to Queue
    messageId = await sendMessage(event);
  } catch (err) {
    // 3. Error handling (final touch)
    const errorResponse: APIGatewayProxyResult = await handleError(err);
    return errorResponse;
  }

  // 4. Response
  return {
    statusCode: 200,
    body: JSON.stringify({
      messageId,
      message: 'success'
    }),
  }
}

Goals:

Translate errors into error responses (ie VerifySignatureError, AwsSqsServiceError, Error, ...)

Manage the logging for the errors

To sum it up, our goal should be to create a utility function - handleError that takes care of all the errors returns the error response.

These goals will act as a guide for us as we develop the module and it will also inform us about what we should test for.

Now we have a good understanding of what we need, let’s start coding this out.

Error handler utility function

1. Create the utility file

touch src/utils/handle-error.ts

2. Add logic to handle the error

Within our error handler, we really have two “known” errors.

The errors are:

VerifySignatureError - When the signature verification fails
AwsSqsServiceError - Errors related to SQS

Then we have the default case which is any other error we haven’t handled in our code or any uncaught errors.

// src/utils/handle-error.ts

import {
  APIGatewayProxyResult
} from 'aws-lambda';

import {
  CommonError,
  AwsSqsServiceError,
  VerifySignatureError,
} from '@app/errors';

import { ServiceError } from '@app/types';

export default function handleError(
  error: ServiceError | Error
) : APIGatewayProxyResult {
  const response : any = {
    statusCode: 500,
    body: {
      message: 'Something went wrong',
      errors: []
    },
  };
  switch (error.constructor.name) {
    // Authentication failure or signature mis-match
    case VerifySignatureError.name:
      response.statusCode = 401;
      break;
    // SQS error - server error
    case AwsSqsServiceError.name:
      response.body.message = error.message;
      break;
    default:
      response.body.message = error.message;
      break;
  }
  response.body.errors.push(error.message);
  response.body = JSON.stringify(response.body);
  return response;
}

💡 Note: You may have noticed that VerifySignatureError is considered a 4xx error.
This is intended. It is a client side error related to an unauthorized request. So, that means we should reject it.

3. Add error tracking id to the response

By default, API gateway will return the request id (awsRequestId) but that is only in the header so we will also include it in our response data for better client debugging.

This is also useful if we ever want to show that identifier on the UI.

// src/utils/handle-error.ts

import {
  APIGatewayProxyResult
} from 'aws-lambda';

import {
  CommonError,
  AwsSqsServiceError,
  VerifySignatureError,
} from '@app/errors';

import asyncLocalStorage from '@app/utils/async-local-storage';

import { ServiceError } from '@app/types';

export default function handleError(
  error: ServiceError | Error
) : APIGatewayProxyResult {
  const requestId: string = asyncLocalStorage.getStore().get('awsRequestId');
  const response : any = {
    statusCode: 500,
    headers: {
      'X-Error-Tracking-Id': requestId,
    },
    body: {
      errorTrackingId: requestId,
      message: 'Something went wrong',
      errors: []
    },
  };
  switch (error.constructor.name) {
    // Authentication failure or signature mis-match
    case VerifySignatureError.name:
      response.statusCode = 401;
      break;
    // SQS error - server error
    case AwsSqsServiceError.name:
      response.body.message = error.message;
      break;
    default:
      response.body.message = error.message;
      break;
  }
  response.body.errors.push(error.message);
  response.body = JSON.stringify(response.body);
  return response;
}

4. Include logging

import {
  APIGatewayProxyResult
} from 'aws-lambda';

import {
  CommonError,
  AwsSqsServiceError,
  VerifySignatureError,
} from '@app/errors';

import logger from '@app/services/logger';
import asyncLocalStorage from '@app/utils/async-local-storage';

import { ServiceError } from '@app/types';

function handleLog(
  error: ServiceError | Error,
  response: APIGatewayProxyResult,
) {
  let logDetails: any = { message: error.message };
  if (error instanceof CommonError) {
    logDetails.operation = error.operation;
    logDetails.context = error.context;
  }
  logDetails.clientResponse = response;
  logger.error(logDetails);
}

export default function handleError(
  error: ServiceError | Error
) : APIGatewayProxyResult {
  const requestId: string = asyncLocalStorage.getStore().get('awsRequestId');
  const response : any = {
    statusCode: 500,
    body: {
      errorTrackingId: requestId,
      message: 'Something went wrong',
      errors: []
    },
  };
  switch (error.constructor.name) {
    // Authentication failure or signature mis-match
    case VerifySignatureError.name:
      response.statusCode = 401;
      break;
    // SQS error - server error
    case AwsSqsServiceError.name:
      break;
    default:
      break;
  }
  response.body.message = error.message;
  response.body.errors.push(error.message);
  response.body = JSON.stringify(response.body);
  // Logging
  handleLog(error, response);
  return response;
}

5. Export util in index file

// src/utils/index.ts

export { default as verifySignature } from './verify-signature';
export { default as getSqsMessage } from './get-sqs-message';
export { default as handleError } from './handle-error';

6. Add tests for our module

Create new file:

touch src/utils/__tests__/handle-error.test.ts

Add tests:

What we are testing for:

✅ The correct Error responses are being generated
✅ The correct status code are being returned
✅ The correct logging details are being added to the logger

// src/utils/__tests__/handle-error.test.ts

import { handleError } from '@app/utils';

import logger from '@app/services/logger';

import {
  AwsSqsServiceError,
  VerifySignatureError
} from '@app/errors';

const mockRequestId = 'test-123';

jest.mock('@app/services/logger', () => ({
  __esModule: true,
  default: {
    info: jest.fn(),
    error: jest.fn(),
  }
}));

jest.mock('@app/utils/async-local-storage', () => ({
  __esModule: true,
  default: ({
    getStore: () => new Map()
      .set('awsRequestId', mockRequestId),
  })
}));

describe('utils/handle-error', () => {
  beforeEach(() => {
    jest.resetAllMocks();
    jest.clearAllMocks();
  });

  // Unhandled error
  describe('case: unhandled error', () => {
    it('should return 500 as status code', () => {
      const error = new Error('test');
      const errorResponse = handleError(error);
      expect(errorResponse.statusCode).toBe(500);
    });
    it('should return the correct body response (JSON.stringify)', () => {
      const error = new Error('test');
      const errorResponse = handleError(error);
      expect(errorResponse.body).toBe(JSON.stringify({
        errorTrackingId: mockRequestId,
        message: error.message,
        errors: [error.message]
      }));
    });
  });

  // VerifySignatureError
  describe('case: VerifySignatureError', () => {
    it('should return 401 as status code (unauthorized)', () => {
      const error = new VerifySignatureError('test');
      const errorResponse = handleError(error);
      expect(errorResponse.statusCode).toBe(401);
    });
    it('should return the correct body response (JSON.stringify)', () => {
      const error = new VerifySignatureError('test');
      const errorResponse = handleError(error);
      expect(errorResponse.body).toBe(JSON.stringify({
        errorTrackingId: mockRequestId,
        message: error.message,
        errors: [error.message]
      }));
    });
    it('should include the context in the logger', () => {
      const error = new VerifySignatureError('test')
        .setContext({
          a: 'a',
          b: 'b'
        });
      const response = handleError(error);
      expect(logger.error).toBeCalledWith({
        message: error.message,
        context: {
          a: 'a',
          b: 'b'
        },
        clientResponse: response,
        operation: ''
      });
    });
    it('should include the operation in the logger', () => {
      const error = new VerifySignatureError('test')
        .setContext({
          a: 'a',
          b: 'b'
        })
        .setOperation('utils/verify-signature');
      const response = handleError(error);
      expect(logger.error).toBeCalledWith({
        message: error.message,
        context: {
          a: 'a',
          b: 'b'
        },
        clientResponse: response,
        operation: 'utils/verify-signature'
      });
    });
  });

  // AwsSqsServiceError
  describe('case: AwsSqsServiceError', () => {
    it('should return 500 as status code', () => {
      const error = new AwsSqsServiceError('test');
      const errorResponse = handleError(error);
      expect(errorResponse.statusCode).toBe(500);
    });
    it('should return the correct body response (JSON.stringify)', () => {
      const error = new AwsSqsServiceError('test');
      const errorResponse = handleError(error);
      expect(errorResponse.body).toBe(JSON.stringify({
        errorTrackingId: mockRequestId,
        message: error.message,
        errors: [error.message]
      }));
    });
    it('should include the context in the logger', () => {
      const error = new AwsSqsServiceError('test')
        .setContext({
          a: 'a',
          b: 'b'
        });
      const response = handleError(error);
      expect(logger.error).toBeCalledWith({
        message: error.message,
        context: {
          a: 'a',
          b: 'b'
        },
        clientResponse: response,
        operation: ''
      });
    });
    it('should include the operation in the logger', () => {
      const error = new AwsSqsServiceError('test')
        .setContext({
          a: 'a',
          b: 'b'
        })
        .setOperation('services/sqs-service');
      const response = handleError(error);
      expect(logger.error).toBeCalledWith({
        message: error.message,
        context: {
          a: 'a',
          b: 'b'
        },
        clientResponse: response,
        operation: 'services/sqs-service'
      });
    });
  });
});

Final integration

Now that we have our handleError utility, let’s integrate that into our lambda function — It should look almost identical to our scaffold we saw earlier!

1. Integrate `handleError` into lambda function

// src/index.ts
import {
  APIGatewayProxyEvent,
  APIGatewayProxyResult
} from 'aws-lambda';

import { sendMessage } from '@app/services/sqs-service';
import {
  handleError,
  verifySignature,
} from '@app/utils';

import { captureRequestContext } from '@app/utils/async-local-storage';

export const handler = async(
  event: APIGatewayProxyEvent
): Promise<APIGatewayProxyResult> => {
  captureRequestContext(event);
  let messageId: string = '';
  try {
    // 1. Verify Signature
    verifySignature(event);
    // 2. Add to Queue
    messageId = await sendMessage(event);
  } catch (err: any) {
    // 3. Error handling (final touch)
    const errorResponse: APIGatewayProxyResult = handleError(err);
    return errorResponse;
  }

  // 4. Response
  return {
    statusCode: 200,
    body: JSON.stringify({
      messageId,
      message: 'success'
    }),
  }
}

2. Tweak the Lambda tests

Let’s also ensure that the lambda function returns the intended responses when the error occurs.

Adding these tests would be closer to integration tests, and somewhat duplicates our unit tests for handleError.

However, it is still useful when we are making changes or refactoring the code.

What we are testing for:

✅ 401 error response is returned on VerifySignatureError
✅ 500 error response is returned on AwsSqsServiceError

The tests:

import { handler } from './index';
import { mockEvent } from '@app/__mocks__';
import { sqs } from '@app/services/sqs-service';
import {
  AwsSqsServiceError,
  VerifySignatureError,
} from '@app/errors';
import verifySignature from '@app/utils/verify-signature';

jest.mock('@app/utils/verify-signature', () => ({
  __esModule: true,
  default: jest.fn(),
}));

jest.mock('@app/services/logger', () => ({
  __esModule: true,
  default: {
    info: jest.fn(),
    error: jest.fn(),
  }
}));

jest.mock('aws-xray-sdk', () => ({
  __esModule: true,
  default: {
    captureAWSClient: jest.fn(),
  }
}));

jest.mock('aws-sdk', () => {
  const SqsMethods = {
    sendMessage: jest.fn().mockReturnThis(),
    promise: jest.fn(),
  }
  const mockAwsSdk = {
    SQS: jest.fn(() => SqsMethods),
  };
  return mockAwsSdk;
});

jest.mock('@app/config', () => ({
  __esModule: true,
  default: {
    webhook: {
      signature: {
        secret: 'test123',
        algo: 'sha256',
        header: 'x-hub-signature-256'
      }
    },
    queue: {
      sqs: {
        url: 'queueUrl'
      }
    },
  }
}));

describe('lambda.handler', () => {
  afterEach(() => {
    jest.resetAllMocks();
    jest.clearAllMocks();
  });
  it('should respond with 200 with the default message', async() => {
    (sqs.sendMessage as jest.Mock).mockReturnValue({
      promise: jest.fn(),
    });
    // @ts-ignore
    await expect(handler(mockEvent))
      .resolves
      .toEqual({
        statusCode: 200,
        body: expect.stringMatching(
          'success'
        )
      });
  });
  it('should respond with 500 if sqs fails', async() => {
    (sqs.sendMessage as jest.Mock).mockReturnValue({
      promise: jest.fn().mockImplementation(() => {
        throw new AwsSqsServiceError('test');
      }),
    });
    // @ts-ignore
    await expect(handler(mockEvent))
      .resolves
      .toEqual({
        statusCode: 500,
        body: JSON.stringify({
          errorTrackingId: mockEvent.requestContext.requestId,
          message: 'failed to sendMessage',
          errors: [
            'failed to sendMessage'
          ]
        })
      });
  });
  it('should response with 401 if signature verification fails', async() => {
    verifySignature
      .mockImplementation(() => {
        throw new VerifySignatureError('test')
      });
    // @ts-ignore
    await expect(handler(mockEvent))
      .resolves
      .toEqual({
        statusCode: 401,
        body: JSON.stringify({
          errorTrackingId: mockEvent.requestContext.requestId,
          message: 'test',
          errors: [
            'test'
          ]
        })
      });
  });
});

Conclusion

Here is the link to the completed version of this module - build-a-webhook-microservice-error-handling.

To recap, simple and testable code start at the first step, with the design. Meaning, it must not be an after thought.

It must be considered and designed.

The reason why most code is not testable is because it was not designed to be tested.

Considerations around testing came only after the code was written or it was not considered at all.

The same applies to simplicity.

We saw this process in action, by designing, implementing and testing the error handler for the Lambda function.

During the design phase, we were very clear on the goals, and on the boundaries these “units” should occupy. Then we carefully written tests to reinforce them.

So, the biggest take away from this article would be that the design & planning of your code is just as important as as code and the tests you write.

And... That’s it. I hope this was helpful or you learned something new!

If you got value out of this, please help spread the word by sharing this article with a friend or co-worker 🙏❤️! (Thanks!)

Enjoy the content ?

Then consider signing up to get notified when new content arrives!