From a46815e5358a92864042eec6d2f13e5a2d06ed48 Mon Sep 17 00:00:00 2001
From: Ilya Muradyan <ilya.muradyan@jetbrains.com>
Date: Mon, 3 Nov 2025 14:35:19 +0100
Subject: [PATCH] [tests] KTNB-1206, KTNB-1222, KTNB-1204 Fix waiting for
 execution to finish

GitOrigin-RevId: 854f643d323e633e0ca5d37d4c18c60532938bd7
---
 .../notebooks/NotebookEditorComponent.kt      | 25 +++++++++++--------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/platform/remote-driver/test-sdk/src/com/intellij/driver/sdk/ui/components/notebooks/NotebookEditorComponent.kt b/platform/remote-driver/test-sdk/src/com/intellij/driver/sdk/ui/components/notebooks/NotebookEditorComponent.kt
index b05b471db487..9a29405881ae 100644
--- a/platform/remote-driver/test-sdk/src/com/intellij/driver/sdk/ui/components/notebooks/NotebookEditorComponent.kt
+++ b/platform/remote-driver/test-sdk/src/com/intellij/driver/sdk/ui/components/notebooks/NotebookEditorComponent.kt
@@ -157,13 +157,13 @@ class NotebookEditorUiComponent(private val data: ComponentData) : JEditorUiComp
     }
   }
 
-  fun runCellAndWaitExecuted(timeout: Duration = 30.seconds): Unit = step("Executing cell") {
+  fun runCellAndWaitExecuted(
+    timeout: Duration = 30.seconds,
+    expectedExecutionCount: Int = 1
+  ): Unit = step("Executing cell") {
     runCell()
     waitFor(timeout = timeout) {
-      // TODO: what if the cell we ran doesn't have an execution label yet, and we are waiting for the previous one?
-      notebookCellExecutionInfos.isNotEmpty() && notebookCellExecutionInfos.last().getParent().x {
-        contains(byAttribute("defaulticon", "greenCheckmark.svg"))
-      }.present()
+      areAllExecutionsFinishedSuccessfully(expectedExecutionCount)
     }
   }
 
@@ -192,14 +192,19 @@ class NotebookEditorUiComponent(private val data: ComponentData) : JEditorUiComp
     waitFor(timeout = timeout) {
       // TODO: what if we have some cells that were executed before, and their checkmarks are still there,
       //  while new execution labels are not yet created?
-      val infos = notebookCellExecutionInfos
+      areAllExecutionsFinishedSuccessfully(notebookCellEditors.size)
+    }
+  }
 
-      infos.isNotEmpty()
-      && infos.size == notebookCellEditors.size
-      && infos.all {
+  private fun areAllExecutionsFinishedSuccessfully(
+    expectedExecutionCount: Int
+  ): Boolean {
+    val infos = notebookCellExecutionInfos
+    return infos.isNotEmpty()
+           && infos.size == expectedExecutionCount
+           && infos.all {
         it.getParent().x { contains(byAttribute("defaulticon", "greenCheckmark.svg")) }.present()
       }
-    }
   }
 
   /**