diff --git a/platform/platform-api/api-dump-unreviewed.txt b/platform/platform-api/api-dump-unreviewed.txt
index 58784adeb520..1db2271d7da5 100644
--- a/platform/platform-api/api-dump-unreviewed.txt
+++ b/platform/platform-api/api-dump-unreviewed.txt
@@ -4447,6 +4447,7 @@ f:com.intellij.openapi.ui.TestDialogManager
 - s:getTestImplementation():com.intellij.openapi.ui.TestDialog
 - s:getTestInputImplementation():com.intellij.openapi.ui.TestInputDialog
 - s:setTestDialog(com.intellij.openapi.ui.TestDialog):com.intellij.openapi.ui.TestDialog
+- s:setTestDialog(com.intellij.openapi.ui.TestDialog,com.intellij.openapi.Disposable):com.intellij.openapi.ui.TestDialog
 - s:setTestInputDialog(com.intellij.openapi.ui.TestInputDialog):com.intellij.openapi.ui.TestInputDialog
 com.intellij.openapi.ui.TestInputDialog
 - sf:DEFAULT:com.intellij.openapi.ui.TestInputDialog
diff --git a/platform/platform-api/src/com/intellij/openapi/ui/TestDialogManager.java b/platform/platform-api/src/com/intellij/openapi/ui/TestDialogManager.java
index 56be21e4aef7..a1072448494b 100644
--- a/platform/platform-api/src/com/intellij/openapi/ui/TestDialogManager.java
+++ b/platform/platform-api/src/com/intellij/openapi/ui/TestDialogManager.java
@@ -1,9 +1,11 @@
 // Copyright 2000-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package com.intellij.openapi.ui;
 
+import com.intellij.openapi.Disposable;
 import com.intellij.openapi.application.Application;
 import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.diagnostic.Logger;
+import com.intellij.openapi.util.Disposer;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.annotations.TestOnly;
@@ -20,6 +22,13 @@ public final class TestDialogManager {
     return oldValue;
   }
 
+  @TestOnly
+  public static TestDialog setTestDialog(@Nullable TestDialog newValue, @NotNull Disposable disposable) {
+    TestDialog oldValue = setTestDialog(newValue);
+    Disposer.register(disposable, () -> implementation = oldValue);
+    return oldValue;
+  }
+
   public static @NotNull TestDialog getTestImplementation() {
     TestDialog result = implementation;
     return result == null ? TestDialog.DEFAULT : result;