move the resetting of sCU to the overridden function

enzymejs · Sep 26, 2017 · 89e2354 · 89e2354
1 parent d5373a3
commit 89e2354
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/packages/enzyme/src/ShallowWrapper.js b/packages/enzyme/src/ShallowWrapper.js
@@ -389,6 +389,7 @@ class ShallowWrapper {
           originalShouldComponentUpdate = instance.shouldComponentUpdate;
           instance.shouldComponentUpdate = (...args) => {
             shouldRender = originalShouldComponentUpdate.apply(instance, args);
+            instance.shouldComponentUpdate = originalShouldComponentUpdate;
             return shouldRender;
           };
         }
@@ -402,9 +403,6 @@ class ShallowWrapper {
         ) {
           instance.componentDidUpdate(prevProps, prevState, prevContext);
         }
-        if (originalShouldComponentUpdate) {
-          instance.shouldComponentUpdate = originalShouldComponentUpdate;
-        }
         this.update();
       });
     });