Scratch tests: check output after repl stop because error output is printed only after repl process termination

This should fix flacky tests

idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ScratchExecutor.kt

@@ -23,8 +23,11 @@ import com.intellij.openapi.editor.event.DocumentEvent
 import com.intellij.openapi.editor.event.DocumentListener
 import com.intellij.openapi.fileEditor.FileDocumentManager
 import com.intellij.psi.PsiManager
+import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.idea.scratch.output.ScratchOutput
 import org.jetbrains.kotlin.idea.scratch.output.ScratchOutputHandler
+import java.util.concurrent.Semaphore
+import java.util.concurrent.TimeUnit
 
 abstract class ScratchExecutor(protected val file: ScratchFile) {
     abstract fun execute()
@@ -176,4 +179,14 @@ abstract class SequentialScratchExecutor(file: ScratchFile) : ScratchExecutor(fi
     private fun wasExpressionExecuted(index: Int): Boolean {
         return index <= lastExecuted
     }
+
+    @TestOnly
+    fun stopAndWait() {
+        val lock = Semaphore(1)
+        lock.acquire()
+        stopExecution {
+            lock.release()
+        }
+        lock.tryAcquire(2, TimeUnit.SECONDS)
+    }
 }

idea/tests/org/jetbrains/kotlin/idea/scratch/AbstractScratchRunActionTest.kt

@@ -212,7 +212,7 @@ abstract class AbstractScratchRunActionTest : FileEditorManagerTestCase() {
         action.actionPerformed(e)
     }
 
-    protected fun waitUntilScratchFinishes() {
+    protected fun waitUntilScratchFinishes(shouldStopRepl: Boolean = true) {
         UIUtil.dispatchAllInvocationEvents()
 
         val start = System.currentTimeMillis()
@@ -221,6 +221,18 @@ abstract class AbstractScratchRunActionTest : FileEditorManagerTestCase() {
             Thread.sleep(100)
         }
 
+        if (shouldStopRepl) stopReplProcess()
+
+        UIUtil.dispatchAllInvocationEvents()
+    }
+
+    protected fun stopReplProcess() {
+        if (myFixture.file != null) {
+            val (_, scratchPanel) = getEditorWithScratchPanel(myManager, myFixture.file.virtualFile)
+                ?: error("Couldn't find scratch panel")
+            scratchPanel.scratchFile.replScratchExecutor?.stopAndWait()
+        }
+
         UIUtil.dispatchAllInvocationEvents()
     }
 
@@ -259,12 +271,6 @@ abstract class AbstractScratchRunActionTest : FileEditorManagerTestCase() {
     }
 
     override fun tearDown() {
-        if (myFixture.file != null) {
-            val (_, scratchPanel) = getEditorWithScratchPanel(myManager, myFixture.file.virtualFile)
-                ?: error("Couldn't find scratch panel")
-            scratchPanel.scratchFile.replScratchExecutor?.stop()
-        }
-
         super.tearDown()
 
         VfsRootAccess.disallowRootAccess(KotlinTestUtils.getHomeDirectory())

idea/tests/org/jetbrains/kotlin/idea/scratch/SequentialScratchExecutorTest.kt

@@ -45,13 +45,13 @@ class SequentialScratchExecutorTest : AbstractScratchRunActionTest() {
 
         try {
             launchScratch()
-            waitUntilScratchFinishes()
+            waitUntilScratchFinishes(shouldStopRepl = false)
 
             for ((text, expected) in expression) {
                 typeAndCheckOutput(text, expected)
             }
         } finally {
-            scratchPanel.scratchFile.replScratchExecutor?.stop()
+            stopReplProcess()
         }
     }
 
@@ -61,7 +61,7 @@ class SequentialScratchExecutorTest : AbstractScratchRunActionTest() {
         myFixture.type(text)
 
         launchAction(RunScratchFromHereAction())
-        waitUntilScratchFinishes()
+        waitUntilScratchFinishes(shouldStopRepl = false)
 
         val inlayAfter = getInlays().filterNot { inlaysBefore.contains(it) }