fix(api-stream): Add throwOnError support in undici.stream

lib/api/api-stream.js

@@ -1,10 +1,10 @@
 'use strict'
 
-const { finished } = require('stream')
+const { finished, PassThrough } = require('stream')
 const {
   InvalidArgumentError,
   InvalidReturnValueError,
-  RequestAbortedError
+  RequestAbortedError, ResponseStatusCodeError
 } = require('../core/errors')
 const util = require('../core/util')
 const { AsyncResource } = require('async_hooks')
@@ -16,7 +16,7 @@ class StreamHandler extends AsyncResource {
       throw new InvalidArgumentError('invalid opts')
     }
 
-    const { signal, method, opaque, body, onInfo, responseHeaders } = opts
+    const { signal, method, opaque, body, onInfo, responseHeaders, throwOnError } = opts
 
     try {
       if (typeof callback !== 'function') {
@@ -57,6 +57,7 @@ class StreamHandler extends AsyncResource {
     this.trailers = null
     this.body = body
     this.onInfo = onInfo || null
+    this.throwOnError = throwOnError
 
     if (util.isStream(body)) {
       body.on('error', (err) => {
@@ -76,8 +77,8 @@ class StreamHandler extends AsyncResource {
     this.context = context
   }
 
-  onHeaders (statusCode, rawHeaders, resume) {
-    const { factory, opaque, context } = this
+  onHeaders (statusCode, rawHeaders, resume, statusMessage) {
+    const { factory, opaque, context, callback } = this
 
     if (statusCode < 200) {
       if (this.onInfo) {
@@ -89,6 +90,25 @@ class StreamHandler extends AsyncResource {
 
     this.factory = null
     const headers = this.responseHeaders === 'raw' ? util.parseRawHeaders(rawHeaders) : util.parseHeaders(rawHeaders)
+
+    if (this.throwOnError && statusCode >= 400) {
+      const body = new PassThrough()
+      const chunks = []
+
+      body.on('data', (chunk) => {
+        chunks.push(chunk)
+      })
+
+      body.on('end', () => {
+        const payload = Buffer.concat(chunks).toString()
+
+        this.runInAsyncScope(callback, null, new ResponseStatusCodeError(`Response status code ${statusCode}${statusMessage ? `: ${statusMessage}` : ''}`, statusCode, headers, payload))
+      })
+
+      this.res = body
+      return
+    }
+
     const res = this.runInAsyncScope(factory, null, {
       statusCode,
       headers,

lib/core/util.js

@@ -282,9 +282,9 @@ function isDisturbed (body) {
     stream.isDisturbed
       ? stream.isDisturbed(body) || body[kBodyUsed] // TODO (fix): Why is body[kBodyUsed] needed?
       : body[kBodyUsed] ||
-        body.readableDidRead ||
-        (body._readableState && body._readableState.dataEmitted) ||
-        isReadableAborted(body)
+      body.readableDidRead ||
+      (body._readableState && body._readableState.dataEmitted) ||
+      isReadableAborted(body)
   ))
 }
 

test/client-stream.js

@@ -786,3 +786,44 @@ test('stream legacy needDrain', (t) => {
     })
   })
 })
+
+test('steam throw if statusCode >= 400', (t) => {
+  t.plan(3)
+
+  const expectedBodyContent = 'expected_body_content'
+  const expectedStatus = 400
+
+  const server = createServer((req, res) => {
+    res.writeHead(expectedStatus, { 'Content-Type': 'text/plain' })
+    res.end(expectedBodyContent)
+  })
+  t.teardown(server.close.bind(server))
+
+  server.listen(0, async () => {
+    const client = new Client(`http://localhost:${server.address().port}`)
+    t.teardown(client.close.bind(client))
+
+    let written = false
+
+    try {
+      await client.stream({
+        path: '/',
+        method: 'GET',
+        throwOnError: true,
+        opaque: new PassThrough()
+      }, ({ opaque }) => {
+        opaque.on('data', (chunk) => {
+          written = true
+        })
+
+        return opaque
+      })
+
+      t.fail('No Error')
+    } catch (e) {
+      t.false(written)
+      t.equal(e.status, expectedStatus)
+      t.equal(e.body, expectedBodyContent)
+    }
+  })
+})