[HMR] Better handle react-hot-loader (#8897)

src/styles/withStyles.js

@@ -163,12 +163,24 @@ const withStyles = (
         this.attach(this.theme);
 
         // Rerender the component so the underlying component gets the theme update.
+        // By theme update we mean receiving and applying the new class names.
         this.setState({}, () => {
           this.detach(oldTheme);
         });
       });
     }
 
+    componentWillReceiveProps() {
+      // react-hot-loader specific logic
+      if (this.stylesCreatorSaved === stylesCreator || process.env.NODE_ENV === 'production') {
+        return;
+      }
+
+      this.detach(this.theme);
+      this.stylesCreatorSaved = stylesCreator;
+      this.attach(this.theme);
+    }
+
     componentWillUnmount() {
       this.detach(this.theme);
 

src/styles/withStyles.spec.js

@@ -206,4 +206,25 @@ describe('withStyles', () => {
       });
     });
   });
+
+  describe('HMR with same state', () => {
+    it('should take the new stylesCreator into account', () => {
+      const styles1 = { root: { padding: 1 } };
+      const StyledComponent1 = withStyles(styles1, { name: 'MuiTextField' })(Empty);
+      const wrapper = shallow(<StyledComponent1 />);
+
+      const styles2 = { root: { padding: 2 } };
+      const StyledComponent2 = withStyles(styles2, { name: 'MuiTextField' })(Empty);
+
+      // Simulate react-hot-loader behavior
+      wrapper.instance().componentWillReceiveProps = // $FlowExpectedError
+        StyledComponent2.prototype.componentWillReceiveProps;
+
+      const classes1 = wrapper.props().classes.root;
+      wrapper.setProps({});
+      const classes2 = wrapper.props().classes.root;
+
+      assert.notStrictEqual(classes1, classes2, 'should generate new classes');
+    });
+  });
 });

src/transitions/Collapse.spec.js

@@ -40,6 +40,7 @@ describe('<Collapse />', () => {
     assert.strictEqual(child.childAt(0).is('div'), true, 'should be a div');
     assert.strictEqual(
       child
+        .childAt(0)
         .childAt(0)
         .children()
         .type(),