Allow app component to be wrapped with custom enhancer when rendering (#4762)

* Add support for custom App and Component enhancers

* Add ctx.renderPage test

* Add tests for single enhancer function

* Cleanup renderPage options check

* Cleanup

* Add comment about backwards compatibility for renderPage

* Add more test cases

server/render.js

@@ -81,12 +81,23 @@ async function doRender (req, res, pathname, query, {
   // the response might be finshed on the getinitialprops call
   if (isResSent(res)) return
 
-  const renderPage = (enhancer = Page => Page) => {
-    const app = createElement(App, {
-      Component: enhancer(Component),
-      router,
-      ...props
-    })
+  const renderPage = (options = Page => Page) => {
+    let EnhancedApp = App
+    let EnhancedComponent = Component
+
+    // For backwards compatibility
+    if (typeof options === 'function') {
+      EnhancedComponent = options(Component)
+    } else if (typeof options === 'object') {
+      if (options.enhanceApp) {
+        EnhancedApp = options.enhanceApp(App)
+      }
+      if (options.enhanceComponent) {
+        EnhancedComponent = options.enhanceComponent(Component)
+      }
+    }
+
+    const app = createElement(EnhancedApp, { Component: EnhancedComponent, router, ...props })
 
     const render = staticMarkup ? renderToStaticMarkup : renderToString
 

test/integration/app-document/pages/_document.js

@@ -2,8 +2,26 @@ import Document, { Head, Main, NextScript } from 'next/document'
 
 export default class MyDocument extends Document {
   static async getInitialProps (ctx) {
-    const initialProps = await Document.getInitialProps(ctx)
-    return { ...initialProps, customProperty: 'Hello Document' }
+    let options
+
+    const enhanceComponent = Component => (props) => <div><span id='render-page-enhance-component'>RENDERED</span><Component {...props} /></div>
+    const enhanceApp = Component => (props) => <div><span id='render-page-enhance-app'>RENDERED</span><Component {...props} /></div>
+
+    if (ctx.query.withEnhancer) {
+      options = enhanceComponent
+    } else if (ctx.query.withEnhanceComponent || ctx.query.withEnhanceApp) {
+      options = {}
+      if (ctx.query.withEnhanceComponent) {
+        options.enhanceComponent = enhanceComponent
+      }
+      if (ctx.query.withEnhanceApp) {
+        options.enhanceApp = enhanceApp
+      }
+    }
+
+    const result = ctx.renderPage(options)
+
+    return { ...result, customProperty: 'Hello Document' }
   }
 
   render () {

test/integration/app-document/test/rendering.js

@@ -34,6 +34,31 @@ export default function ({ app }, suiteName, render, fetch) {
         })
         expect(noncesAdded).toBe(true)
       })
+
+      test('It renders ctx.renderPage with enhancer correctly', async () => {
+        const $ = await get$('/?withEnhancer=true')
+        const nonce = 'RENDERED'
+        expect($('#render-page-enhance-component').text().includes(nonce)).toBe(true)
+      })
+
+      test('It renders ctx.renderPage with enhanceComponent correctly', async () => {
+        const $ = await get$('/?withEnhanceComponent=true')
+        const nonce = 'RENDERED'
+        expect($('#render-page-enhance-component').text().includes(nonce)).toBe(true)
+      })
+
+      test('It renders ctx.renderPage with enhanceApp correctly', async () => {
+        const $ = await get$('/?withEnhanceApp=true')
+        const nonce = 'RENDERED'
+        expect($('#render-page-enhance-app').text().includes(nonce)).toBe(true)
+      })
+
+      test('It renders ctx.renderPage with enhanceApp and enhanceComponent correctly', async () => {
+        const $ = await get$('/?withEnhanceComponent=true&withEnhanceApp=true')
+        const nonce = 'RENDERED'
+        expect($('#render-page-enhance-app').text().includes(nonce)).toBe(true)
+        expect($('#render-page-enhance-component').text().includes(nonce)).toBe(true)
+      })
     })
 
     describe('_app', () => {