bug: correctly implement multiple reads

Author: kessler

index.js

@@ -1,41 +1,61 @@
+const debug = require('debug')('aws-api-read-stream')
 const { Readable, finished } = require('stream')
 
+const STOPPED = 'stopped'
+const READING = 'reading'
+const DONE = 'done'
+
 class AWSApiReadStream extends Readable {
 	constructor(fn, opts, nextToken) {
 		super({ ...opts, objectMode: true })
 
 		this._fn = fn
 		this._nextToken = nextToken
+		this._state = STOPPED
 	}
 
-	async _read(size) {
-		let shouldContinue = false
-		do {
-			try {
-				const res = await this._fn(this._nextToken)
+	_read(size) {
+		if (this._state === READING) return
+		debug('_read', size, 'nextToken', this._nextToken)
+		this._execApiCall()
+	}
 
-				if (!res) {
-					this.push(null)
-					return
-				}
+	async _execApiCall() {
+		this._state = READING
+		try {
+			const res = await this._fn(this._nextToken)
+			
+			if (!res) {
+				this._apiExecutionDone()
+				return
+			}
 
-				if (this._isInBufferMode()) {
-					this._buffer.push(res)
-				}
+			if (this._isInBufferMode()) {
+				this._buffer.push(res)
+			}
 
-				shouldContinue = this.push(res)
-				this._nextToken = res.NextToken || res.NextContinuationToken
+			this._nextToken = res.NextToken || res.NextContinuationToken
 
-				if (!this._nextToken) {
-					this.push(null)
-					return
-				}
+			if (!this._nextToken) {
+				this._apiExecutionDone()
+				return
+			}
 
-			} catch (e) {
-				this.destroy(e)
+			if (!this.push(res)) {
+				this._state = STOPPED
 				return
 			}
-		} while (shouldContinue)
+
+			this._execApiCall()
+		} catch (e) {
+			this.destroy(e)
+		}
+	}
+
+	_apiExecutionDone() {
+		this._state = DONE
+		this._nextToken = undefined
+		this.push(null)
 	}
 
 	stop() {

package-lock.json

@@ -4,7 +4,7 @@
   "description": "turns an aws api call into a readable stream",
   "main": "index.js",
   "scripts": {
-    "test": "npx ava"
+    "test": "env DEBUG=aws-api-read-stream npx ava"
   },
   "repository": {
     "type": "git",
@@ -25,5 +25,8 @@
   "devDependencies": {
     "ava": "^3.5.2",
     "aws-sdk": "^2.649.0"
+  },
+  "dependencies": {
+    "debug": "^4.1.1"
   }
 }

package.json

@@ -1,5 +1,6 @@
 const test = require('ava')
 const AWSApiReadStream = require('./index')
+const tokens = 'abcdefghij'.split('')
 
 test('AWSApiReadStream', async t => {
 	const testapi = new TestAPI()
@@ -9,35 +10,35 @@ test('AWSApiReadStream', async t => {
 
 	for await (const { data, NextToken } of stream) {
 		t.is(data, count)
-
-		if (count < 10) {
-			t.is(NextToken, count + 1)
-
-			t.is(testapi.nextTokenSpy, count + 1)
-
-		} else {
-			t.is(NextToken, undefined)
-		}
-
 		count++
 	}
+
+	t.deepEqual(testapi.spy, [
+		{ nextToken: undefined, data: 0 },
+		{ nextToken: 'a', data: 1 },
+		{ nextToken: 'b', data: 2 },
+		{ nextToken: 'c', data: 3 },
+		{ nextToken: 'd', data: 4 },
+		{ nextToken: 'e', data: 5 },
+		{ nextToken: 'f', data: 6 },
+		{ nextToken: 'g', data: 7 },
+		{ nextToken: 'h', data: 8 },
+		{ nextToken: 'i', data: 9 },
+		{ nextToken: 'j', data: 10 }
+	])
 })
 
 test('AWSApiReadStream - initialize with existing token', async t => {
 	const testapi = new TestAPI()
-
-	let count = 0
-
-	const stream = AWSApiReadStream.from(nextToken => testapi.anAPICall(nextToken), { nextToken: 5 })
+	const stream = AWSApiReadStream.from(nextToken => testapi.anAPICall(nextToken), { nextToken: 'f' })
 
 	// can't think of a better way to test this other than check internals...
-	t.is(stream._nextToken, 5)
+	t.is(stream._nextToken, 'f')
 })
 
 test('AWSApiReadStream - returning null or undefined will stop execution', async t => {
 	const testapi = new TestAPI()
-	const stream = AWSApiReadStream.from(nextToken => nextToken === 2 ? null : testapi.anAPICall(nextToken))
-
+	const stream = AWSApiReadStream.from(nextToken => nextToken === 'b' ? null : testapi.anAPICall(nextToken))
 
 	const results = []
 	for await (const { data, NextToken } of stream) {
@@ -50,19 +51,19 @@ test.skip('AWSApiReadStream - backpressure', async t => {
 
 })
 
-
 class TestAPI {
 	constructor() {
 		this._counter = 0
+		this.spy = []
 	}
 
 	anAPICall(nextToken) {
-		this.nextTokenSpy = nextToken
+		this.spy.push({ nextToken, data: this._counter })
 		return new Promise(res => {
 			setImmediate(() => {
 				res({
-					data: this._counter,
-					NextToken: this._counter < 10 ? ++this._counter : undefined
+					NextToken: tokens[this._counter],
+					data: this._counter++
 				})
 			})
 		})