serverless/test/integration/aws/websocket.test.js

'use strict'

const WebSocket = require('ws')
const { expect } = require('chai')
const awsRequest = require('@serverless/test/aws-request')
const CloudFormationService = require('aws-sdk').CloudFormation
const log = require('log').get('serverless:test')
const wait = require('timers-ext/promise/sleep')
const fixtures = require('../../fixtures/programmatic')

const { confirmCloudWatchLogs } = require('../../utils/misc')
const { deployService, removeService } = require('../../utils/integration')
const {
  createApi,
  deleteApi,
  getRoutes,
  createStage,
  deleteStage,
} = require('../../utils/websocket')

describe('AWS - API Gateway Websocket Integration Test', function () {
  this.timeout(1000 * 60 * 10) // Involves time-taking deploys
  let stackName
  let serviceName
  let serviceDir
  let updateConfig
  const stage = 'dev'

  before(async () => {
    const serviceData = await fixtures.setup('websocket')
    ;({ servicePath: serviceDir, updateConfig } = serviceData)
    serviceName = serviceData.serviceConfig.service
    stackName = `${serviceName}-${stage}`
    return deployService(serviceDir)
  })

  after(() => removeService(serviceDir))

  async function getWebSocketServerUrl() {
    const result = await awsRequest(CloudFormationService, 'describeStacks', {
      StackName: stackName,
    })
    const webSocketServerUrl = result.Stacks[0].Outputs.find(
      (output) => output.OutputKey === 'ServiceEndpointWebsocket',
    ).OutputValue
    return webSocketServerUrl
  }

  describe('Two-Way Setup', () => {
    let timeoutId
    after(() => clearTimeout(timeoutId))

    it('should expose a websocket route that can reply to a message', async () => {
      const webSocketServerUrl = await getWebSocketServerUrl()

      return new Promise((resolve, reject) => {
        const ws = new WebSocket(webSocketServerUrl)
        reject = ((promiseReject) => (error) => {
          promiseReject(error)
          try {
            ws.close()
          } catch (closeError) {
            // safe to ignore
          }
        })(reject)

        const sendMessage = () => {
          log.debug("Sending message to 'hello' route")
          ws.send(JSON.stringify({ action: 'hello', name: 'serverless' }))
          timeoutId = setTimeout(sendMessage, 1000)
        }

        ws.on('error', reject)
        ws.on('open', sendMessage)

        ws.on('close', resolve)

        ws.on('message', (event) => {
          clearTimeout(timeoutId)
          try {
            log.debug(`Received WebSocket message: ${event}`)
            expect(event).to.equal('Hello, serverless')
          } finally {
            ws.close()
          }
        })
      }).finally(() => clearTimeout(timeoutId))
    })
  })

  describe('Minimal Setup', () => {
    it('should expose an accessible websocket endpoint', async () => {
      const webSocketServerUrl = await getWebSocketServerUrl()

      log.debug(`WebSocket Server URL ${webSocketServerUrl}`)
      expect(webSocketServerUrl).to.match(
        /wss:\/\/.+\.execute-api\..+\.amazonaws\.com.+/,
      )
      return new Promise((resolve, reject) => {
        const ws = new WebSocket(webSocketServerUrl)
        let isRejected = false
        reject = ((promiseReject) => (error) => {
          isRejected = true
          promiseReject(error)
          try {
            ws.close()
          } catch (closeError) {
            // safe to ignore
          }
        })(reject)
        ws.on('error', reject)
        ws.on('open', () => {
          confirmCloudWatchLogs(`/aws/websocket/${stackName}`, () => {
            if (isRejected) throw new Error('Stop propagation')
            ws.send('test message')
            return wait(500)
          }).then((events) => {
            expect(events.length > 0).to.equal(true)
            ws.close()
          }, reject)
        })

        ws.on('close', resolve)

        ws.on('message', (event) => {
          log.debug('Unexpected WebSocket message', event)
          reject(new Error('Unexpected message'))
        })
      })
    })

    // NOTE: this test should  be at the very end because we're using an external REST API here
    describe('when using an existing websocket API', () => {
      let websocketApiId
      before(async () => {
        // create an external websocket API
        const externalWebsocketApiName = `${stage}-${serviceName}-ext-api`
        const wsApiMeta = await createApi(externalWebsocketApiName)
        websocketApiId = wsApiMeta.ApiId
        await createStage(websocketApiId, 'dev')
        await updateConfig({
          provider: {
            apiGateway: { websocketApiId },
          },
        })
        return deployService(serviceDir)
      })

      after(async () => {
        // NOTE: deleting the references to the old, external websocket API
        await updateConfig({
          provider: {
            apiGateway: { websocketApiId: null },
          },
        })
        // NOTE: we need to delete the stage before deleting the stack
        // otherwise CF will refuse to delete the deployment because a stage refers to that
        await deleteStage(websocketApiId, 'dev')
        // NOTE: deploying once again to get the stack into the original state
        await deployService(serviceDir)
        log.debug('Deleting external websocket API...')
        await deleteApi(websocketApiId)
      })

      it('should add the routes to the referenced API', async () => {
        const routes = await getRoutes(websocketApiId)
        expect(routes.length).to.equal(4)
      })
    })
  })
})