fix(gatsby-react-router-scroll): scroll restoration for layout components (#26861) (#31079)

* add generic type to useScrollRestoration hook

this is required to match the type of `useRef` which requires the target
element type to be specified.

The hook can be used like the following:
```typescript
const MyComponent: React.FunctionComponent = () => {
  const scrollRestorationProps = useScrollRestoration<HTMLDivElement>(`some-key`)
  return (
    <div
      {...scrollRestorationProps}
      style={{ overflow: `auto` }}
    >
      Test
    </div>
  )
}
```

Fixes: #26458

* add tests for useScrollRestoration hook

* fix useScrollResotration to update position on location change

Previously, the scroll position was only updated when the using
component was re-rendered. This did not work when the hook is used in
a wrapping Layout component, becaus its persitent over location changes.
Adding the location key to useEffect will cause the scroll position is
updated every time the key changes.

* lint/ts fixes

Co-authored-by: Vladimir Razuvaev <[email protected]>
(cherry picked from commit f57efab)

Co-authored-by: Michael van Engelshoven <[email protected]>

Author: GatsbyJS Bot

packages/gatsby-react-router-scroll/src/__tests__/use-scroll-restoration.tsx

@@ -0,0 +1,171 @@
+import React from "react"
+import {
+  LocationProvider,
+  History,
+  createMemorySource,
+  createHistory,
+} from "@reach/router"
+import { render, fireEvent } from "@testing-library/react"
+import { useScrollRestoration } from "../use-scroll-restoration"
+import { ScrollHandler } from "../scroll-handler"
+import { SessionStorage } from "../session-storage"
+
+const TRUE = (): boolean => true
+
+const Fixture: React.FunctionComponent = () => {
+  const scrollRestorationProps = useScrollRestoration<HTMLDivElement>(`test`)
+  return (
+    <div
+      {...scrollRestorationProps}
+      style={{ overflow: `auto` }}
+      data-testid="scrollfixture"
+    >
+      Test
+    </div>
+  )
+}
+
+describe(`useScrollRestoration`, () => {
+  let history: History
+  const session = new SessionStorage()
+  let htmlElementPrototype: HTMLElement
+  let fakedScrollTo = false
+
+  beforeAll(() => {
+    const wrapper = render(<div>hello</div>)
+    htmlElementPrototype = wrapper.container.constructor.prototype
+
+    // jsdom doesn't support .scrollTo(), lets fix this temporarily
+    if (typeof htmlElementPrototype.scrollTo === `undefined`) {
+      htmlElementPrototype.scrollTo = function scrollTo(
+        optionsOrX?: ScrollToOptions | number,
+        y?: number
+      ): void {
+        if (typeof optionsOrX === `number`) {
+          this.scrollLeft = optionsOrX
+        }
+        if (typeof y === `number`) {
+          this.scrollTop = y
+        }
+      }
+      fakedScrollTo = true
+    }
+  })
+
+  beforeEach(() => {
+    history = createHistory(createMemorySource(`/`))
+    sessionStorage.clear()
+  })
+
+  afterAll(() => {
+    if (fakedScrollTo && htmlElementPrototype.scrollTo) {
+      // @ts-ignore
+      delete htmlElementPrototype.scrollTo
+    }
+  })
+
+  it(`stores current scroll position in storage`, () => {
+    const wrapper = render(
+      <LocationProvider history={history}>
+        <ScrollHandler
+          navigate={history.navigate}
+          location={history.location}
+          shouldUpdateScroll={TRUE}
+        >
+          <Fixture />
+        </ScrollHandler>
+      </LocationProvider>
+    )
+
+    fireEvent.scroll(wrapper.getByTestId(`scrollfixture`), {
+      target: { scrollTop: 123 },
+    })
+
+    expect(session.read(history.location, `test`)).toBe(123)
+  })
+
+  it(`scrolls to stored offset on render`, () => {
+    session.save(history.location, `test`, 684)
+
+    const wrapper = render(
+      <LocationProvider history={history}>
+        <ScrollHandler
+          navigate={history.navigate}
+          location={history.location}
+          shouldUpdateScroll={TRUE}
+        >
+          <Fixture />
+        </ScrollHandler>
+      </LocationProvider>
+    )
+
+    expect(wrapper.getByTestId(`scrollfixture`)).toHaveProperty(
+      `scrollTop`,
+      684
+    )
+  })
+
+  it(`scrolls to 0 on render when session has no entry`, () => {
+    const wrapper = render(
+      <LocationProvider history={history}>
+        <ScrollHandler
+          navigate={history.navigate}
+          location={history.location}
+          shouldUpdateScroll={TRUE}
+        >
+          <Fixture />
+        </ScrollHandler>
+      </LocationProvider>
+    )
+
+    expect(wrapper.getByTestId(`scrollfixture`)).toHaveProperty(`scrollTop`, 0)
+  })
+
+  it(`updates scroll position on location change`, async () => {
+    const wrapper = render(
+      <LocationProvider history={history}>
+        <ScrollHandler
+          navigate={history.navigate}
+          location={history.location}
+          shouldUpdateScroll={TRUE}
+        >
+          <Fixture />
+        </ScrollHandler>
+      </LocationProvider>
+    )
+
+    fireEvent.scroll(wrapper.getByTestId(`scrollfixture`), {
+      target: { scrollTop: 356 },
+    })
+
+    await history.navigate(`/another-location`)
+
+    expect(wrapper.getByTestId(`scrollfixture`)).toHaveProperty(`scrollTop`, 0)
+  })
+
+  it(`restores scroll position when navigating back`, async () => {
+    const wrapper = render(
+      <LocationProvider history={history}>
+        <ScrollHandler
+          navigate={history.navigate}
+          location={history.location}
+          shouldUpdateScroll={TRUE}
+        >
+          <Fixture />
+        </ScrollHandler>
+      </LocationProvider>
+    )
+
+    fireEvent.scroll(wrapper.getByTestId(`scrollfixture`), {
+      target: { scrollTop: 356 },
+    })
+
+    await history.navigate(`/another-location`)
+    await history.navigate(-1)
+
+    expect(wrapper.getByTestId(`scrollfixture`)).toHaveProperty(
+      `scrollTop`,
+      356
+    )
+  })
+})

packages/gatsby-react-router-scroll/src/use-scroll-restoration.ts

@@ -2,24 +2,24 @@ import { ScrollContext } from "./scroll-handler"
 import { useRef, useContext, useLayoutEffect, MutableRefObject } from "react"
 import { useLocation } from "@gatsbyjs/reach-router"
 
-interface IScrollRestorationProps {
-  ref: MutableRefObject<HTMLElement | undefined>
+interface IScrollRestorationProps<T extends HTMLElement> {
+  ref: MutableRefObject<T | null>
   onScroll(): void
 }
 
-export function useScrollRestoration(
+export function useScrollRestoration<T extends HTMLElement>(
   identifier: string
-): IScrollRestorationProps {
+): IScrollRestorationProps<T> {
   const location = useLocation()
   const state = useContext(ScrollContext)
-  const ref = useRef<HTMLElement>()
+  const ref = useRef<T>(null)
 
   useLayoutEffect((): void => {
     if (ref.current) {
       const position = state.read(location, identifier)
       ref.current.scrollTo(0, position || 0)
     }
-  }, [])
+  }, [location.key])
 
   return {
     ref,