expose abort function

googleapis · Jul 16, 2015 · 99a2476 · 99a2476
1 parent 04b3e6c
commit 99a2476
Show file tree

Hide file tree

Showing 2 changed files with 94 additions and 6 deletions.
diff --git a/index.js b/index.js
@@ -41,6 +41,7 @@ function retryRequest(requestOpts, opts, callback) {
 
   if (streamMode) {
     retryStream = through();
+    retryStream.abort = resetStreams;
   }
 
   makeRequest();
@@ -51,8 +52,11 @@ function retryRequest(requestOpts, opts, callback) {
 
   function resetStreams() {
     cacheStream = null;
-    requestStream.abort();
-    requestStream.destroy();
+
+    if (requestStream) {
+      requestStream.abort();
+      requestStream.destroy();
+    }
   }
 
   function makeRequest() {
@@ -62,7 +66,7 @@ function retryRequest(requestOpts, opts, callback) {
       cacheStream = new StreamCache();
       requestStream = opts.request(requestOpts);
 
-      streamForward(requestStream)
+      streamForward(requestStream, { events: ['response', 'complete'] })
         .on('error', onResponse)
         .on('response', onResponse.bind(null, null))
         .pipe(cacheStream);
@@ -97,8 +101,8 @@ function retryRequest(requestOpts, opts, callback) {
     // No more attempts need to be made, just continue on.
     if (streamMode) {
       streamForward(cacheStream)
-        .on('error', resetStreams)
-        .pipe(retryStream);
+        .pipe(retryStream)
+        .on('error', resetStreams);
     } else {
       callback(err, response, body);
     }

diff --git a/test.js b/test.js
@@ -2,6 +2,7 @@
 
 var assert = require('assert');
 var retryRequest = require('./index.js');
+var through = require('through2');
 
 describe('retry-request', function () {
   var URI_404 = 'http://yahoo.com/theblahstore';
@@ -29,6 +30,89 @@ describe('retry-request', function () {
           done();
         });
     });
+
+    it('exposes an `abort` fuction to match request', function (done) {
+      var retryStream = retryRequest(URI_NON_EXISTENT);
+
+      retryStream.on('error', function () {
+        assert.equal(typeof retryStream.abort, 'function');
+        done();
+      });
+    });
+
+    it('works on the last attempt', function (done) {
+      var numAborts = 0;
+      var numAttempts = 0;
+
+      var opts = {
+        request: function () {
+          numAttempts++;
+
+          var fakeRequestStream = through();
+          fakeRequestStream.abort = function () {
+            numAborts++;
+          };
+
+          var response = numAttempts < 3 ? { statusCode: 503 } : { statusCode: 200 };
+          setImmediate(function () {
+            fakeRequestStream.emit('response', response);
+
+            setImmediate(function () {
+              // They all emit 'complete', but the user's stream should only get
+              // the last one.
+              fakeRequestStream.emit('complete', numAttempts);
+            });
+          });
+
+          return fakeRequestStream;
+        }
+      };
+
+      retryRequest(URI_404, opts)
+        .on('error', done)
+        .on('complete', function (numAttempts) {
+          assert.strictEqual(numAborts, 2);
+          assert.deepEqual(numAttempts, 3);
+          done();
+        });
+    });
+
+    it('never succeeds', function (done) {
+      var numAborts = 0;
+      var numAttempts = 0;
+
+      var opts = {
+        request: function () {
+          numAttempts++;
+
+          var fakeRequestStream = through();
+          fakeRequestStream.abort = function () {
+            numAborts++;
+          };
+
+          var response = { statusCode: 503 };
+          setImmediate(function () {
+            fakeRequestStream.emit('response', response);
+
+            setImmediate(function () {
+              // They all emit 'complete', but the user's stream should only get
+              // the last one.
+              fakeRequestStream.emit('complete', numAttempts);
+            });
+          });
+
+          return fakeRequestStream;
+        }
+      };
+
+      retryRequest(URI_404, opts)
+        .on('error', done)
+        .on('complete', function (numAttempts) {
+          assert.strictEqual(numAborts, 2);
+          assert.strictEqual(numAttempts, 3);
+          done();
+        });
+    });
   });
 
   describe('callbacks', function () {
@@ -59,7 +143,7 @@ describe('retry-request', function () {
       var shouldRetryFnCalled = false;
 
       var opts = {
-        retries: 1, // so that our retry function is only called once and
+        retries: 1, // so that our retry function is only called once
 
         shouldRetryFn: function () {
           shouldRetryFnCalled = true;