From 40c4ffadf888ecf8ab043e290f7d4a043764b38b Mon Sep 17 00:00:00 2001
From: Alex Zherdev <alex.zherdev@gmail.com>
Date: Sat, 23 Jun 2018 19:36:13 -0700
Subject: [PATCH] Account for UNSAFE_ method in no-will-update-set-state

Resolves #1844
---
 lib/rules/no-will-update-set-state.js       |  6 ++-
 lib/util/makeNoMethodSetStateRule.js        | 20 +++++++---
 tests/lib/rules/no-will-update-set-state.js | 41 +++++++++++++++++++++
 3 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/lib/rules/no-will-update-set-state.js b/lib/rules/no-will-update-set-state.js
index 51f44ebe59..22bb633114 100644
--- a/lib/rules/no-will-update-set-state.js
+++ b/lib/rules/no-will-update-set-state.js
@@ -5,5 +5,9 @@
 'use strict';
 
 const makeNoMethodSetStateRule = require('../util/makeNoMethodSetStateRule');
+const versionUtil = require('../util/version');
 
-module.exports = makeNoMethodSetStateRule('componentWillUpdate');
+module.exports = makeNoMethodSetStateRule(
+  'componentWillUpdate',
+  context => versionUtil.testReactVersion(context, '16.3.0')
+);
diff --git a/lib/util/makeNoMethodSetStateRule.js b/lib/util/makeNoMethodSetStateRule.js
index 2b497581e9..dcc9b09555 100644
--- a/lib/util/makeNoMethodSetStateRule.js
+++ b/lib/util/makeNoMethodSetStateRule.js
@@ -10,7 +10,7 @@ const docsUrl = require('./docsUrl');
 // Rule Definition
 // ------------------------------------------------------------------------------
 
-function makeNoMethodSetStateRule(methodName) {
+function makeNoMethodSetStateRule(methodName, shouldCheckUnsafeCb) {
   return {
     meta: {
       docs: {
@@ -28,6 +28,15 @@ function makeNoMethodSetStateRule(methodName) {
     create: function(context) {
       const mode = context.options[0] || 'allow-in-func';
 
+      function nameMatches(name) {
+        let matches = name === methodName;
+        if (typeof shouldCheckUnsafeCb === 'function' && shouldCheckUnsafeCb(context)) {
+          matches = matches || name === `UNSAFE_${methodName}`;
+        }
+
+        return matches;
+      }
+
       // --------------------------------------------------------------------------
       // Public
       // --------------------------------------------------------------------------
@@ -46,19 +55,20 @@ function makeNoMethodSetStateRule(methodName) {
           const ancestors = context.getAncestors(callee).reverse();
           let depth = 0;
           for (let i = 0, j = ancestors.length; i < j; i++) {
-            if (/Function(Expression|Declaration)$/.test(ancestors[i].type)) {
+            const ancestor = ancestors[i];
+            if (/Function(Expression|Declaration)$/.test(ancestor.type)) {
               depth++;
             }
             if (
-              (ancestors[i].type !== 'Property' && ancestors[i].type !== 'MethodDefinition') ||
-              ancestors[i].key.name !== methodName ||
+              (ancestor.type !== 'Property' && ancestor.type !== 'MethodDefinition') ||
+              !nameMatches(ancestor.key.name) ||
               (mode !== 'disallow-in-func' && depth > 1)
             ) {
               continue;
             }
             context.report({
               node: callee,
-              message: `Do not use setState in ${methodName}`
+              message: `Do not use setState in ${ancestor.key.name}`
             });
             break;
           }
diff --git a/tests/lib/rules/no-will-update-set-state.js b/tests/lib/rules/no-will-update-set-state.js
index 25afca0b42..e131c0755c 100644
--- a/tests/lib/rules/no-will-update-set-state.js
+++ b/tests/lib/rules/no-will-update-set-state.js
@@ -77,6 +77,18 @@ ruleTester.run('no-will-update-set-state', rule, {
       });
     `,
     parser: 'babel-eslint'
+  }, {
+    code: `
+      class Hello extends React.Component {
+        UNSAFE_componentWillUpdate() {
+          this.setState({
+            data: data
+          });
+        }
+      }
+    `,
+    parser: 'babel-eslint',
+    settings: {react: {version: '16.2.0'}}
   }],
 
   invalid: [{
@@ -225,5 +237,34 @@ ruleTester.run('no-will-update-set-state', rule, {
     errors: [{
       message: 'Do not use setState in componentWillUpdate'
     }]
+  }, {
+    code: `
+      class Hello extends React.Component {
+        UNSAFE_componentWillUpdate() {
+          this.setState({
+            data: data
+          });
+        }
+      }
+    `,
+    parser: 'babel-eslint',
+    settings: {react: {version: '16.3.0'}},
+    errors: [{
+      message: 'Do not use setState in UNSAFE_componentWillUpdate'
+    }]
+  }, {
+    code: `
+      var Hello = createReactClass({
+        UNSAFE_componentWillUpdate: function() {
+          this.setState({
+            data: data
+          });
+        }
+      });
+    `,
+    settings: {react: {version: '16.3.0'}},
+    errors: [{
+      message: 'Do not use setState in UNSAFE_componentWillUpdate'
+    }]
   }]
 });