Added more test cases

packages/react-reconciler/src/ReactFiberScheduler.js

@@ -1293,6 +1293,12 @@ function commitRootImpl(root) {
     // This usually means we've finished all the work, but it can also happen
     // when something gets downprioritized during render, like a hidden tree.
     root.lastPendingTime = firstPendingTimeBeforeCommit;
+
+    if (enableSchedulerTracing) {
+      if (firstPendingTimeBeforeCommit !== NoWork) {
+        restorePendingSchedulers(root, root.lastPendingTime);
+      }
+    }
   }
 
   if (root === workInProgressRoot) {
@@ -2183,7 +2189,6 @@ export function restorePendingSchedulers(
   if (!enableUpdateSchedulerTracking) {
     return;
   }
-
   const pendingSchedulersMap = root.pendingSchedulersMap;
   let schedulers = pendingSchedulersMap.get(expirationTime);
   if (schedulers == null) {
@@ -2322,6 +2327,10 @@ function startWorkOnPendingInteraction(root, expirationTime) {
       }
     });
 
+    // Store the current set of interactions on the FiberRoot for a few reasons:
+    // We can re-use it in hot functions like renderRoot() without having to
+    // recalculate it. This also provides DevTools with a way to access it when
+    // the onCommitRoot() hook is called.
     root.memoizedSchedulers = memoizedSchedulers;
   }
 

packages/react-reconciler/src/ReactFiberUnwindWork.js

@@ -209,6 +209,11 @@ function throwException(
   // Its effect list is no longer valid.
   sourceFiber.firstEffect = sourceFiber.lastEffect = null;
 
+  if (enableUpdateSchedulerTracking) {
+    // If we have pending work still, restore the original schedulers
+    restorePendingSchedulers(root, renderExpirationTime);
+  }
+
   if (
     value !== null &&
     typeof value === 'object' &&
@@ -331,10 +336,6 @@ function throwException(
           if (enableSchedulerTracing) {
             retry = Schedule_tracing_wrap(retry);
           }
-          if (enableUpdateSchedulerTracking) {
-            // If we have pending work still, restore the original schedulers
-            restorePendingSchedulers(root, renderExpirationTime);
-          }
           thenable.then(retry, retry);
         }
         workInProgress.effectTag |= ShouldCapture;