Repository: koajs/koa
File: __tests__/response/flushHeaders.test.js (JavaScript)

Code before:
    const PassThrough = require('stream').PassThrough
    const app = new Koa()

    app.use(ctx => {
      ctx.type = 'json'
      ctx.status = 200
      ctx.headers.Link = '</css/mycss.css>; as=style; rel=preload, <https://img.craftflair.com>; rel=preconnect; crossorigin'
      const stream = ctx.body = new PassThrough()
      ctx.flushHeaders()
      setTimeout(() => {
        stream.end(JSON.stringify({ message: 'hello!' }))
      })
    })

    const server = app.listen()
    await once(server, 'listening')
    const port = server.address().port

    try {
      const req = http.request({ port })
      req.end()

      const [res] = await once(req, 'response')
      const onData = () => { throw new Error('boom') }
      res.on('data', onData)
      res.removeListener('data', onData)
      res.destroy()
    } finally {
      server.close()
    }
  })

  it('should catch stream error', async () => {
    const PassThrough = require('stream').PassThrough
    const app = new Koa()
    app.once('error', err => {
      assert(err.message === 'mock error')
    })

    app.use(ctx => {
      ctx.type = 'json'
      ctx.status = 200
      ctx.headers.Link = '</css/mycss.css>; as=style; rel=preload, <https://img.craftflair.com>; rel=preconnect; crossorigin'
      ctx.length = 20
      ctx.flushHeaders()
      const stream = ctx.body = new PassThrough()

      setTimeout(() => {
        stream.emit('error', new Error('mock error'))
        stream.end()
      })

Review comment (suggestion):
this test is not valid since the listener is added and removed and the test is complete on the same tick - the logic only happens on next tick. we can keep it for now but it's a bad test

Code after:
    const PassThrough = require('stream').PassThrough
    const app = new Koa()
    let headersFlushed = false
    let dataReceived = false

    app.use(ctx => {
      ctx.type = 'json'
      ctx.status = 200
      ctx.headers.Link = '</css/mycss.css>; as=style; rel=preload, <https://img.craftflair.com>; rel=preconnect; crossorigin'
      const stream = ctx.body = new PassThrough()
      ctx.flushHeaders()
      headersFlushed = true
      setTimeout(() => {
        stream.end(JSON.stringify({ message: 'hello!' }))
      }, 10)
    })

    const server = app.listen()
    await once(server, 'listening')
    const port = server.address().port

    try {
      const req = http.request({ port })
      req.end()

      const [res] = await once(req, 'response')
      assert(headersFlushed, 'Headers should be flushed')

      const dataPromise = new Promise(resolve => {
        res.once('data', chunk => {
          dataReceived = true
          resolve(chunk)
        })
      })

      // Wait for data with a timeout
      const timeoutPromise = new Promise((resolve, reject) =>
        setTimeout(() => reject(new Error('Timeout waiting for data')), 100)
      )

      await Promise.race([dataPromise, timeoutPromise])
      res.destroy()
      assert(dataReceived, 'Data should be received after headers')
    } finally {
      server.close()
    }
  })

  it('should catch stream error', async () => {
    const PassThrough = require('stream').PassThrough
    const app = new Koa()