From b2d0aa643c0199b75bb4acb516c508e6ef102c18 Mon Sep 17 00:00:00 2001
From: Aaron Corley <aaryn101@gmail.com>
Date: Mon, 3 Aug 2020 10:57:32 -0600
Subject: [PATCH] fix(useSortBy): fix the bug that causes flatRows to be
 incorrect when sorting and using subRows (#2443)

Co-authored-by: Aaron Corley <aaron.corley@sirsidynix.com>
Co-authored-by: Tanner Linsley <tannerlinsley@gmail.com>
---
 src/plugin-hooks/tests/useSortBy.test.js | 84 ++++++++++++++++++++----
 src/plugin-hooks/useSortBy.js            |  5 +-
 2 files changed, 74 insertions(+), 15 deletions(-)

diff --git a/src/plugin-hooks/tests/useSortBy.test.js b/src/plugin-hooks/tests/useSortBy.test.js
index e856c28514..109674c72c 100644
--- a/src/plugin-hooks/tests/useSortBy.test.js
+++ b/src/plugin-hooks/tests/useSortBy.test.js
@@ -28,20 +28,32 @@ const data = [
     status: 'Complicated',
     progress: 10,
   },
+  {
+    firstName: 'john',
+    lastName: 'buggyman',
+    age: 52,
+    visits: 24,
+    status: 'Married',
+    progress: 17,
+    subRows: [
+      {
+        firstName: 'winston',
+        lastName: 'buggyman',
+        age: 18,
+        visits: 200,
+        status: 'Single',
+        progress: 10,
+      },
+    ],
+  },
 ]
 
 const defaultColumn = {
   Cell: ({ value, column: { id } }) => `${id}: ${value}`,
 }
 
-function Table({ columns, data }) {
-  const {
-    getTableProps,
-    getTableBodyProps,
-    headerGroups,
-    rows,
-    prepareRow,
-  } = useTable(
+function Table({ columns, data, useTableRef }) {
+  const instance = useTable(
     {
       columns,
       data,
@@ -50,6 +62,18 @@ function Table({ columns, data }) {
     useSortBy
   )
 
+  const {
+    getTableProps,
+    getTableBodyProps,
+    headerGroups,
+    rows,
+    prepareRow,
+  } = instance
+
+  if (useTableRef) {
+    useTableRef.current = instance
+  }
+
   return (
     <table {...getTableProps()}>
       <thead>
@@ -85,7 +109,7 @@ function Table({ columns, data }) {
   )
 }
 
-function App() {
+function App({ useTableRef }) {
   const columns = React.useMemo(
     () => [
       {
@@ -126,7 +150,7 @@ function App() {
     []
   )
 
-  return <Table columns={columns} data={data} />
+  return <Table columns={columns} data={data} useTableRef={useTableRef} />
 }
 
 test('renders a sortable table', () => {
@@ -139,7 +163,12 @@ test('renders a sortable table', () => {
       .queryAllByRole('row')
       .slice(2)
       .map(d => d.children[0].textContent)
-  ).toEqual(['firstName: derek', 'firstName: joe', 'firstName: tanner'])
+  ).toEqual([
+    'firstName: derek',
+    'firstName: joe',
+    'firstName: john',
+    'firstName: tanner',
+  ])
 
   fireEvent.click(rendered.getByText('First Name 🔼0'))
   rendered.getByText('First Name 🔽0')
@@ -148,7 +177,12 @@ test('renders a sortable table', () => {
       .queryAllByRole('row')
       .slice(2)
       .map(d => d.children[0].textContent)
-  ).toEqual(['firstName: tanner', 'firstName: joe', 'firstName: derek'])
+  ).toEqual([
+    'firstName: tanner',
+    'firstName: john',
+    'firstName: joe',
+    'firstName: derek',
+  ])
 
   fireEvent.click(rendered.getByText('Profile Progress'))
   rendered.getByText('Profile Progress 🔼0')
@@ -157,7 +191,12 @@ test('renders a sortable table', () => {
       .queryAllByRole('row')
       .slice(2)
       .map(d => d.children[0].textContent)
-  ).toEqual(['firstName: joe', 'firstName: tanner', 'firstName: derek'])
+  ).toEqual([
+    'firstName: joe',
+    'firstName: john',
+    'firstName: tanner',
+    'firstName: derek',
+  ])
 
   fireEvent.click(rendered.getByText('First Name'), { shiftKey: true })
   rendered.getByText('Profile Progress 🔼0')
@@ -167,5 +206,22 @@ test('renders a sortable table', () => {
       .queryAllByRole('row')
       .slice(2)
       .map(d => d.children[0].textContent)
-  ).toEqual(['firstName: joe', 'firstName: derek', 'firstName: tanner'])
+  ).toEqual([
+    'firstName: joe',
+    'firstName: john',
+    'firstName: derek',
+    'firstName: tanner',
+  ])
+})
+
+test('maintains the integrity of instance.flatRows', () => {
+  const useTableRef = { current: null }
+  const rendered = render(<App useTableRef={useTableRef} />)
+
+  fireEvent.click(rendered.getByText('First Name'))
+  const flatRows = useTableRef.current.flatRows
+  expect(flatRows.length).toBe(5)
+  expect(
+    flatRows.find(r => r.values.firstName === 'winston')
+  ).not.toBeUndefined()
 })
diff --git a/src/plugin-hooks/useSortBy.js b/src/plugin-hooks/useSortBy.js
index d551636706..825fc13c55 100755
--- a/src/plugin-hooks/useSortBy.js
+++ b/src/plugin-hooks/useSortBy.js
@@ -321,7 +321,10 @@ function useInstance(instance) {
       // If there are sub-rows, sort them
       sortedData.forEach(row => {
         sortedFlatRows.push(row)
-        if (!row.subRows || row.subRows.length < 1) {
+        if (!row.subRows) {
+          return
+        } else if (row.subRows.length === 1) {
+          sortedFlatRows.push(row.subRows[0])
           return
         }
         row.subRows = sortData(row.subRows)