move the resetting of sCU to the overridden function

Author: lelandrichardson

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();
       });
     });