diff --git a/platform/remote-servers/target-integration-tests/testSrc/com/intellij/tests/targets/java/CommonJavaTargetTestBase.kt b/platform/remote-servers/target-integration-tests/testSrc/com/intellij/tests/targets/java/CommonJavaTargetTestBase.kt
index de8dbef6fb13..9943c50be886 100644
--- a/platform/remote-servers/target-integration-tests/testSrc/com/intellij/tests/targets/java/CommonJavaTargetTestBase.kt
+++ b/platform/remote-servers/target-integration-tests/testSrc/com/intellij/tests/targets/java/CommonJavaTargetTestBase.kt
@@ -304,14 +304,7 @@ abstract class CommonJavaTargetTestBase(protected val executionMode: ExecutionMo
   }
 
   private fun runWithConnectInUnitTestMode(block: () -> Unit) {
-    val connectInUnitTestModeValue = TestModeFlags.get(SearchForTestsTask.CONNECT_IN_UNIT_TEST_MODE_PROPERTY_KEY)
-    TestModeFlags.set(SearchForTestsTask.CONNECT_IN_UNIT_TEST_MODE_PROPERTY_KEY, true)
-    try {
-      block()
-    }
-    finally {
-      TestModeFlags.set(SearchForTestsTask.CONNECT_IN_UNIT_TEST_MODE_PROPERTY_KEY, connectInUnitTestModeValue)
-    }
+    TestModeFlags.runWithFlag(SearchForTestsTask.CONNECT_IN_UNIT_TEST_MODE_PROPERTY_KEY, true, block)
   }
 
   /**