Added before and after hooks

README.md

@@ -32,11 +32,31 @@ var options = {
   protocol: 'http',                 // optional, default: this.protocol
   host: 'www.risingstack.com'       // optional, default: this.host,
   prerenderToken: ''                // optional or process.env.PRERENDER_TOKEN
+  beforePrerender: function *() {   // optional before hook
+    this.body = 'foo';
+  },
+  afterPrerender: function *() {    // optional after hook
+    this.body = 'bar';
+  }
 };
 
 // Use as middleware
 app.use(prerender(options));
+
 ```
+> When prerendering happens the request will not yield further down the stack but return the rendered page instead.
+
+
+#### Hooks
+
+If `beforePrerender` or `afterPrerender` is set in the options, they will be called
+before and after the prerendering is supposed to happen. If `beforePrerender` sets a
+body, it will be used and the actual prerender server will **not be called** for that
+request.
+
+This makes it possible to reliably tell when a request was or is going to be prerendered and therefore the perfect place
+to handle caching if needed.
+
 
 ## License
 

index.js

@@ -149,31 +149,34 @@ module.exports = function preRenderMiddleware (options) {
       url: this.url
     });
 
-    var body = '';
-
     var renderUrl;
     var preRenderUrl;
-    var response;
 
     // Pre-render generate the site and return
     if (isPreRender) {
       renderUrl = protocol + '://' + host + this.url;
       preRenderUrl = options.prerender + renderUrl;
-      response = yield requestGet({
-        url: preRenderUrl,
-        headers: headers,
-        gzip: true
-      });
-
-      body = response[1] || '';
-
-      yield* next;
+      if (options.beforePrerender) {
+        yield options.beforePrerender.call(this);
+      }
+
+      if (! this.body) {
+        var response = yield requestGet({
+          url: preRenderUrl,
+          headers: headers,
+          gzip: true
+        });
+        this.body = (response[1] || '').toString();
+      }
 
-      this.body = body.toString();
       this.set('X-Prerender', 'true');
+
+      if (options.afterPrerender) {
+        yield options.afterPrerender.call(this);
+      }
     } else {
-      yield* next;
       this.set('X-Prerender', 'false');
+      yield* next;
     }
   };
 };

package.json

@@ -1,6 +1,6 @@
 {
   "name": "koa-prerender",
-  "version": "1.0.3",
+  "version": "2.0.0",
   "description": "",
   "main": "index.js",
   "scripts": {

test.js

@@ -50,4 +50,77 @@ describe('Koa prerender middleware', function() {
 
   });
 
+  describe('when prerenders', function() {
+    describe('calls hooks', function () {
+      it('when beforePrerender is defined', function (done) {
+        var app = koa();
+        app.use(prerender({
+          'beforePrerender': function *() {
+            this.body = 'foo'
+          }
+        }));
+
+        request(app.listen())
+          .get('/?_escaped_fragment_')
+          .expect('X-Prerender', 'true')
+          .expect(200, 'foo', done)
+      });
+
+      it('when afterPrerender is defined', function (done) {
+        var app = koa();
+        app.use(prerender({
+          'afterPrerender': function *() {
+            this.body = 'foo'
+          }
+        }));
+
+        request(app.listen())
+          .get('/?_escaped_fragment_')
+          .expect('X-Prerender', 'true')
+          .expect(200, 'foo', done)
+      })
+    });
+
+    it('does not yield', function (done) {
+      var app = koa();
+      app.use(prerender());
+      app.use(function *() {
+        throw new Error('It yielded.');
+      });
+      request(app.listen())
+        .get('/?_escaped_fragment_')
+        .expect(200, done)
+    });
+  });
+  describe('when does not prerender', function() {
+    it('does not call hooks', function(done) {
+      var app = koa();
+      app.use(prerender({
+        'beforePrerender': function *() {
+          this.body = 'foo'
+        },
+        'afterPrerender': function *() {
+          this.body = 'bar'
+        }
+      }));
+
+      request(app.listen())
+        .get('/')
+        .expect('X-Prerender', 'false')
+        .expect(404, done)
+    });
+
+    it('yields', function(done) {
+      var app = koa();
+      app.use(prerender());
+      app.use(function *() {
+        this.body = 'foo'
+      });
+      request(app.listen())
+        .get('/')
+        .expect('X-Prerender', 'false')
+        .expect(200, 'foo', done)
+    })
+  });
+
 });