diff --git a/src/stats/bw.js b/src/stats/bw.js index 3c0d79c87..7a9d882fc 100644 --- a/src/stats/bw.js +++ b/src/stats/bw.js @@ -2,20 +2,15 @@ const promisify = require('promisify-es6') const streamToValue = require('../utils/stream-to-value') +const { Transform } = require('readable-stream') -const transform = function (res, callback) { - streamToValue(res, (err, data) => { - if (err) { - return callback(err) - } - - callback(null, { - totalIn: data[0].TotalIn, - totalOut: data[0].TotalOut, - rateIn: data[0].RateIn, - rateOut: data[0].RateOut - }) - }) +const transformChunk = (chunk) => { + return { + totalIn: chunk.TotalIn, + totalOut: chunk.TotalOut, + rateIn: chunk.RateIn, + rateOut: chunk.RateOut + } } module.exports = (send) => { @@ -28,6 +23,28 @@ module.exports = (send) => { send.andTransform({ path: 'stats/bw', qs: opts - }, transform, callback) + }, (res, callback) => { + if (!opts.poll) { + // If not polling, just send the result. + return streamToValue(res, (err, data) => { + if (err) { + return callback(err) + } + + callback(null, transformChunk(data[0])) + }) + } + + // If polling, return a readable stream. + const output = new Transform({ + objectMode: true, + transform (chunk, encoding, cb) { + cb(null, transformChunk(chunk)) + } + }) + + res.pipe(output) + callback(null, output) + }, callback) }) }