From 1dcf2c4326f389d6f423f8e42d25dad250dcbaf8 Mon Sep 17 00:00:00 2001
From: Grace Guo <grace.guo@airbnb.com>
Date: Thu, 15 Jun 2017 15:30:13 -0700
Subject: [PATCH] fix local state 'columns' (#2896)

* fix local state 'columns'
* fix method name per code review
---
 .../SqlLab/components/VisualizeModal.jsx      | 12 +++--
 .../sqllab/VisualizeModal_spec.jsx            | 45 ++++++++++++++-----
 2 files changed, 40 insertions(+), 17 deletions(-)

diff --git a/superset/assets/javascripts/SqlLab/components/VisualizeModal.jsx b/superset/assets/javascripts/SqlLab/components/VisualizeModal.jsx
index 7c19645a68..2ebfef9318 100644
--- a/superset/assets/javascripts/SqlLab/components/VisualizeModal.jsx
+++ b/superset/assets/javascripts/SqlLab/components/VisualizeModal.jsx
@@ -41,28 +41,26 @@ class VisualizeModal extends React.PureComponent {
     this.state = {
       chartType: CHART_TYPES[0],
       datasourceName: this.datasourceName(),
-      columns: {},
+      columns: this.getColumnFromProps(),
       hints: [],
     };
   }
   componentDidMount() {
     this.validate();
   }
-  componentWillReceiveProps(nextProps) {
-    this.setStateFromProps(nextProps);
-  }
-  setStateFromProps(props) {
+  getColumnFromProps() {
+    const props = this.props;
     if (!props ||
         !props.query ||
         !props.query.results ||
         !props.query.results.columns) {
-      return;
+      return {};
     }
     const columns = {};
     props.query.results.columns.forEach((col) => {
       columns[col.name] = col;
     });
-    this.setState({ columns });
+    return columns;
   }
   datasourceName() {
     const { query } = this.props;
diff --git a/superset/assets/spec/javascripts/sqllab/VisualizeModal_spec.jsx b/superset/assets/spec/javascripts/sqllab/VisualizeModal_spec.jsx
index b9d36a365d..eab4608889 100644
--- a/superset/assets/spec/javascripts/sqllab/VisualizeModal_spec.jsx
+++ b/superset/assets/spec/javascripts/sqllab/VisualizeModal_spec.jsx
@@ -78,20 +78,45 @@ describe('VisualizeModal', () => {
     expect(wrapper.find(Modal)).to.have.length(1);
   });
 
-  describe('setStateFromProps', () => {
-    const wrapper = getVisualizeModalWrapper();
-    const sampleQuery = queries[0];
-
-    it('should require valid props parameters', () => {
-      const spy = sinon.spy(wrapper.instance(), 'setState');
-      wrapper.instance().setStateFromProps();
-      expect(spy.callCount).to.equal(0);
-      spy.restore();
+  describe('getColumnFromProps', () => {
+    it('should require valid query parameter in props', () => {
+      const emptyQuery = {
+        show: true,
+        query: {},
+      };
+      const wrapper = shallow(<VisualizeModal {...emptyQuery} />, {
+        context: { store },
+      }).dive();
+      expect(wrapper.state().columns).to.deep.equal({});
     });
     it('should set columns state', () => {
-      wrapper.instance().setStateFromProps({ query: sampleQuery });
+      const wrapper = getVisualizeModalWrapper();
       expect(wrapper.state().columns).to.deep.equal(mockColumns);
     });
+    it('should not change columns state when closing Modal', () => {
+      const wrapper = getVisualizeModalWrapper();
+      expect(wrapper.state().columns).to.deep.equal(mockColumns);
+
+      // first change columns state
+      const newColumns = {
+        ds: {
+          is_date: true,
+          is_dim: false,
+          name: 'ds',
+          type: 'STRING',
+        },
+        name: {
+          is_date: false,
+          is_dim: true,
+          name: 'name',
+          type: 'STRING',
+        },
+      };
+      wrapper.instance().setState({ columns: newColumns });
+      // then close Modal
+      wrapper.setProps({ show: false });
+      expect(wrapper.state().columns).to.deep.equal(newColumns);
+    });
   });
 
   describe('datasourceName', () => {