From ed82295049d58b5ced39d91cd18a1fa8318fd11f Mon Sep 17 00:00:00 2001
From: "Dmitry.Yudin" <dmitry.yudin@jetbrains.com>
Date: Thu, 15 Jan 2026 13:41:50 +0100
Subject: [PATCH] test framework [driver]: Fix BackgroundRun.useDriver error
 handling

GitOrigin-RevId: 02d2c64c6296be03b6ff6f728412d2a8c23779f6
---
 .../starter/driver/engine/BackgroundRun.kt    | 24 +++++++++++++++----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/tools/intellij.tools.ide.starter.driver/src/com/intellij/ide/starter/driver/engine/BackgroundRun.kt b/tools/intellij.tools.ide.starter.driver/src/com/intellij/ide/starter/driver/engine/BackgroundRun.kt
index e6b6f5b10ab2..e3f2314be7eb 100644
--- a/tools/intellij.tools.ide.starter.driver/src/com/intellij/ide/starter/driver/engine/BackgroundRun.kt
+++ b/tools/intellij.tools.ide.starter.driver/src/com/intellij/ide/starter/driver/engine/BackgroundRun.kt
@@ -59,11 +59,25 @@ open class BackgroundRun(override val startResult: Deferred<IDEStartResult>, dri
    * The IDE is closed on any exception, or if it doesn't close automatically after the block execution completes.
    */
   open fun <R> useDriver(closeIdeTimeout: Duration = 1.minutes, block: Driver.() -> R): IDEStartResult {
-    val ideStartResult: IDEStartResult
-    try {
-      driver.withContext { block(this) }
-    }
-    finally {
+    lateinit var ideStartResult: IDEStartResult
+    runCatching {
+      driver.withContext {
+        block(this)
+        if (isConnected) takeScreenshot("beforeIdeClosed")
+      }
+    }.onFailure { e ->
+      runCatching {
+        driver.exitApplication()
+        waitFor(
+          timeout = 15.seconds,
+          errorMessage = {
+            forceKill()
+            "Error on exit application via Driver"
+          },
+        ) { !driver.isConnected }
+      }
+      throw e
+    }.onSuccess {
       ideStartResult = driver.waitToClose(closeIdeTimeout)
     }
     return ideStartResult