Setup AB test with Lambda@edge function

Issue

I needed to setup A/B test two different homepage designs. However Amplify doesn't support per-branch split testing like Netlify does.

Finally, I ended up using lambda@edge to set AB test, wrote the code below.

Code

  1. Server checks cookie X-Experiment
  2. If the cookie exists, redirect proper page depending on the cookie's value so that user can see the same page as previous session.
  3. If the cookie doesn't exist, get random variation and set cookie. Then redirect user to proper page.
  4. Cookie should be disappeared after the experiment is done.

js

// Lambda function for `viewer request` hook
const EXPERIMENTS = {
  '/': {
    cookie: 'X-Experiment',
    endDate: new Date(new Date('2023-09-01').getTime() + 1000 * 60 * 60 * 24 * 28), // 28 days
    variations: [
      { id: 'original', percentage: 50, url: '/' },
      { id: 'variant', percentage: 50, url: '/home-b/' },
    ],
  },
  '/en/': {
    cookie: 'X-Experiment',
    endDate: new Date(new Date('2023-09-01').getTime() + 1000 * 60 * 60 * 24 * 28), // 28 days
    variations: [
      { id: 'original', percentage: 50, url: '/en/' },
      { id: 'variant', percentage: 50, url: '/en/home-b/' },
    ],
  }
}

const getRandomVariation = (variations) => {
  const rnd = Math.random() * 100000
  const percent = rnd / 1000
  let result = null
  let acc = 0

  for (let i = 0; i < variations.length; i++) {
    const item = variations[i]
    if (result === null && percent > 100 - item.percentage - acc) {
      result = item
    }
    acc += item.percentage
  }
  return result
}

const parseExperimentCookie = (cookieHeader, experiment) => {
  const cookies = cookieHeader.value.split(';')
  const cookie = cookies.find((cookieValue) => cookieValue.includes(experiment.cookie))
  const value = cookie.split('=')
  return experiment.variations.find((v) => v.id === value[value.length - 1])
}

const findExperimentCookie = (cookies, experiment) => {
  const cookie = cookies.find((cookie) => cookie.value.includes(experiment.cookie))
  return cookie ? parseExperimentCookie(cookie, experiment) : undefined
}
const isCrawler = (userAgent) => {
  const regex =
    /aolbuild|baidu|bingbot|bingpreview|msnbot|duckduckgo|adsbot-google|googlebot|mediapartners-google|teoma|slurp|yandex|bot|crawl|spider|curl|chrome-lighthouse|google page speed insights/g
  return !!userAgent.match(regex)
}

export const handler = async (event) => {
  const request = event.Records[0].cf.request
  const headers = request.headers
  const user_agent = headers?.['user-agent']?.[0]?.value.toLowerCase()

  const experiment = EXPERIMENTS[request.uri]
  if (!experiment || new Date() > experiment.endDate || !user_agent || isCrawler(user_agent)) {
    // Do not process split test
    return request
  }

  headers.cookie = headers.cookie || []

  // Try to find the variation in the experiment cookie.
  const experimentVariation = findExperimentCookie(headers.cookie, experiment)
  // No cookie is found, determine the variation randomly.
  if (!experimentVariation) {
    const variation = getRandomVariation(experiment.variations)
    const url = request.uri + (request.querystring ? `?${request.querystring}` : '')
    const response = {
      status: 302,
      headers: {
        'cache-control': [
          {
            key: 'Cache-Control',
            value: 'no-store',
          },
        ],
        'set-cookie': [
          {
            key: 'Set-Cookie',
            value: `${experiment.cookie}=${variation.id}; Expires=${experiment.endDate.toUTCString()}; path=/`,
          },
        ],
        location: [
          {
            key: 'Location',
            value: url,
          },
        ],
      },
    }
    return response
  }
  // At this point there's always a cookie for this experiment.
  const destinationUri = experimentVariation.url
  // Update the URI and return the request.
  request.uri = destinationUri
  return request
}