From 8bdc527d5822e981ef37bff9a7c3610ba85011cb Mon Sep 17 00:00:00 2001
From: dchamb <dchambers@sapient.com>
Date: Fri, 6 Jan 2017 11:35:32 +0000
Subject: [PATCH] [TableBody] Fix row selection re-render bug

Non-selection related table re-rendering previously caused the user's
selection to be lost because the state was being reset unnecessarily.

I took the opportunity to remove the prescient `TODO` comment that did
correctly predict that this might be a problem.
---
 src/Table/TableBody.js                        | 19 +++---
 .../Table/SelectionObservingTable.js          | 55 ++++++++++++++++
 .../Table/SelectionObservingTable.spec.js     | 62 +++++++++++++++++++
 3 files changed, 127 insertions(+), 9 deletions(-)
 create mode 100644 test/integration/Table/SelectionObservingTable.js
 create mode 100644 test/integration/Table/SelectionObservingTable.spec.js

diff --git a/src/Table/TableBody.js b/src/Table/TableBody.js
index 6cfec3e882ab89..424961ff9939f6 100644
--- a/src/Table/TableBody.js
+++ b/src/Table/TableBody.js
@@ -129,15 +129,16 @@ class TableBody extends Component {
   }
 
   componentWillReceiveProps(nextProps) {
-    if (this.props.allRowsSelected && !nextProps.allRowsSelected) {
-      this.setState({
-        selectedRows: [],
-      });
-      // TODO: should else be conditional, not run any time props other than allRowsSelected change?
-    } else {
-      this.setState({
-        selectedRows: this.calculatePreselectedRows(nextProps),
-      });
+    if (this.props.allRowsSelected !== nextProps.allRowsSelected) {
+      if (!nextProps.allRowsSelected) {
+        this.setState({
+          selectedRows: [],
+        });
+      } else {
+        this.setState({
+          selectedRows: this.calculatePreselectedRows(nextProps),
+        });
+      }
     }
   }
 
diff --git a/test/integration/Table/SelectionObservingTable.js b/test/integration/Table/SelectionObservingTable.js
new file mode 100644
index 00000000000000..0d0749135ae45b
--- /dev/null
+++ b/test/integration/Table/SelectionObservingTable.js
@@ -0,0 +1,55 @@
+import React from 'react';
+import {
+  Table,
+  TableBody,
+  TableHeader,
+  TableHeaderColumn,
+  TableRow,
+  TableRowColumn,
+} from 'src/Table';
+
+const tableData = [
+  {name: 'John Smith'},
+  {name: 'Randal White'},
+  {name: 'Olivier'},
+];
+
+class SelectionObservingTable extends React.Component {
+  constructor(props) {
+    super(props);
+    this.state = {
+      selectionChangeCount: 0,
+    };
+  }
+
+  render() {
+    return (
+      <div>
+        Selection Change Count: {this.state.selectionChangeCount}
+        <Table
+          selectable={true}
+          onRowSelection={() => this.setState({selectionChangeCount: this.state.selectionChangeCount + 1})}
+        >
+          <TableHeader>
+            <TableRow>
+              <TableHeaderColumn>
+                Name
+              </TableHeaderColumn>
+            </TableRow>
+          </TableHeader>
+          <TableBody displayRowCheckbox={true}>
+            {tableData.map( (row, index) => (
+              <TableRow key={index}>
+                <TableRowColumn>
+                  {row.name}
+                </TableRowColumn>
+              </TableRow>
+            ))}
+          </TableBody>
+        </Table>
+      </div>
+    );
+  }
+}
+
+export default SelectionObservingTable;
diff --git a/test/integration/Table/SelectionObservingTable.spec.js b/test/integration/Table/SelectionObservingTable.spec.js
new file mode 100644
index 00000000000000..a06d09d6e7559b
--- /dev/null
+++ b/test/integration/Table/SelectionObservingTable.spec.js
@@ -0,0 +1,62 @@
+/* eslint-env mocha */
+import React from 'react';
+import {mount} from 'enzyme';
+import {assert} from 'chai';
+import getMuiTheme from 'src/styles/getMuiTheme';
+import SelectionObservingTable from './SelectionObservingTable';
+
+describe('<SelectionObservingTable />', () => {
+  it('preserves selected change', () => {
+    const muiTheme = getMuiTheme();
+    const mountWithContext = (node) => mount(node, {
+      context: {muiTheme},
+      childContextTypes: {muiTheme: React.PropTypes.object},
+    });
+
+    const wrapper = mountWithContext(
+      <SelectionObservingTable />
+    );
+
+    assert.deepEqual(
+      wrapper.find('Checkbox').map((checkbox) => checkbox.props().checked),
+      [
+        false,
+        false,
+        false,
+        false,
+      ],
+      'should use the selected property for the initial value'
+    );
+
+    let input;
+    input = wrapper.find('Checkbox').at(1).find('input');
+    input.node.checked = !input.node.checked;
+    input.simulate('change');
+
+    assert.deepEqual(
+      wrapper.find('Checkbox').map((checkbox) => checkbox.props().checked),
+      [
+        false,
+        true,
+        false,
+        false,
+      ],
+      'should preserve initial selection'
+    );
+
+    input = wrapper.find('Checkbox').at(2).find('input');
+    input.node.checked = !input.node.checked;
+    input.simulate('change');
+
+    assert.deepEqual(
+      wrapper.find('Checkbox').map((checkbox) => checkbox.props().checked),
+      [
+        false,
+        false,
+        true,
+        false,
+      ],
+      'should preserve updated selection'
+    );
+  });
+});