Fix document styles with react 18 (#35760)

ijjk · web-flow · commit fe012ff916db · 2022-03-31T00:35:00.000Z
This is a follow-up to #35736 that ensures we use the styles from `_document.getInitialProps` instead of only applying the `styled-jsx` styles. ## Bug - [x] Related issues linked using `fixes #number` - [x] Integration tests added - [ ] Errors have helpful link attached, see `contributing.md` Fixes: #35758 x-ref: #35736
diff --git a/packages/next/server/render.tsx b/packages/next/server/render.tsx
@@ -1431,23 +1431,33 @@ export async function renderToHTML(
         })
       }
 
-      const styles = jsxStyleRegistry.styles()
-      jsxStyleRegistry.flush()
+      const hasDocumentGetInitialProps = !(
+        isServerComponent ||
+        process.browser ||
+        !Document.getInitialProps
+      )
 
-      const documentInitialPropsRes =
-        isServerComponent || process.browser || !Document.getInitialProps
-          ? {}
-          : await documentInitialProps()
+      const documentInitialPropsRes = hasDocumentGetInitialProps
+        ? await documentInitialProps()
+        : {}
       if (documentInitialPropsRes === null) return null
 
+      const { docProps } = (documentInitialPropsRes as any) || {}
       const documentElement = () => {
         if (isServerComponent || process.browser) {
           return (Document as any)()
         }
 
-        const { docProps } = (documentInitialPropsRes as any) || {}
         return <Document {...htmlProps} {...docProps} />
       }
+      let styles
+
+      if (hasDocumentGetInitialProps) {
+        styles = docProps.styles
+      } else {
+        styles = jsxStyleRegistry.styles()
+        jsxStyleRegistry.flush()
+      }
 
       return {
         bodyResult,
diff --git a/test/development/basic/styled-components.test.ts b/test/development/basic/styled-components.test.ts
@@ -2,7 +2,7 @@ import { join } from 'path'
 import webdriver from 'next-webdriver'
 import { createNext, FileRef } from 'e2e-utils'
 import { NextInstance } from 'test/lib/next-modes/base'
-import { fetchViaHTTP } from 'next-test-utils'
+import { fetchViaHTTP, renderViaHTTP } from 'next-test-utils'
 
 describe('styled-components SWC transform', () => {
   let next: NextInstance
@@ -69,4 +69,10 @@ describe('styled-components SWC transform', () => {
       )
     ).toBe('rgb(255, 255, 255)')
   })
+
+  it('should contain styles in initial HTML', async () => {
+    const html = await renderViaHTTP(next.url, '/')
+    expect(html).toContain('background:transparent')
+    expect(html).toContain('color:white')
+  })
 })
