From 0d4b48560f6b51f47591ca5c344144cc4fbb97e3 Mon Sep 17 00:00:00 2001
From: "Kirill.Karnaukhov" <kirill.karnaukhov@jetbrains.com>
Date: Tue, 8 Oct 2024 21:13:39 +0200
Subject: [PATCH] [ai-completion] LLM-11785: fix tests that relied on that
 after `insert` no session is possible

GitOrigin-RevId: e77d1221daf5af09c3ac5154fd54089292f6179a
---
 .../logs/InlineCompletionSessionLogsTest.kt   |  2 ++
 .../logs/TestInlineCompletionLogs.kt          | 19 ++++++++++++++++---
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/logs/InlineCompletionSessionLogsTest.kt b/platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/logs/InlineCompletionSessionLogsTest.kt
index 53da9d65f906..84e566f10584 100644
--- a/platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/logs/InlineCompletionSessionLogsTest.kt
+++ b/platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/logs/InlineCompletionSessionLogsTest.kt
@@ -71,6 +71,7 @@ internal class InlineCompletionSessionLogsTest : InlineCompletionTestCase() {
         callInlineCompletion()
         // now cancel it
         typeChar('a')
+        // it doesn't create the second session log because it hangs infinitely
       }
       singleSessionLog.assertRequestIdPresent()
       singleSessionLog.assertSomeContextLogsPresent()
@@ -94,6 +95,7 @@ internal class InlineCompletionSessionLogsTest : InlineCompletionTestCase() {
         provider.computeNextElement()
         delay()
         typeChar('a')
+        // it doesn't create the second session log because it hangs infinitely
       }
       singleSessionLog.assertRequestIdPresent()
       singleSessionLog.assertSomeContextLogsPresent()
diff --git a/platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/logs/TestInlineCompletionLogs.kt b/platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/logs/TestInlineCompletionLogs.kt
index f042fea50acc..a92daac4e9ed 100644
--- a/platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/logs/TestInlineCompletionLogs.kt
+++ b/platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/logs/TestInlineCompletionLogs.kt
@@ -10,6 +10,7 @@ import com.jetbrains.fus.reporting.model.lion3.LogEvent
 import io.kotest.common.runBlocking
 import org.jetbrains.annotations.ApiStatus
 import kotlin.test.assertEquals
+import kotlin.test.assertNotNull
 import kotlin.test.assertTrue
 
 /**
@@ -42,6 +43,20 @@ object TestInlineCompletionLogs {
    * Also see [noSessionLogs].
    */
   fun InlineCompletionLifecycleTestDSL.singleSessionLog(block: suspend InlineCompletionLifecycleTestDSL.() -> Unit): SingleSessionLog {
+    val sessionsLogs = sessionsLog(block)
+    assertEquals(1, sessionsLogs.size, "There should be exactly session log, but found: $sessionsLogs")
+    return sessionsLogs.single()
+  }
+
+  fun InlineCompletionLifecycleTestDSL.firstSessionLog(block: suspend InlineCompletionLifecycleTestDSL.() -> Unit): SingleSessionLog {
+    val sessionLog = sessionsLog(block).firstOrNull()
+    assertNotNull(sessionLog, "There should be at least one session log, but found nothing.")
+    return sessionLog
+  }
+
+  private fun InlineCompletionLifecycleTestDSL.sessionsLog(
+    block: suspend InlineCompletionLifecycleTestDSL.() -> Unit
+  ): List<SingleSessionLog> {
     val disposable = Disposer.newDisposable()
     val events = try {
       FUCollectorTestCase.collectLogEvents("ML", disposable) {
@@ -53,9 +68,7 @@ object TestInlineCompletionLogs {
     finally {
       Disposer.dispose(disposable)
     }
-    val sessionLogs = getSessionLogs(events)
-    assertEquals(1, sessionLogs.size, "There should be exactly session log, but found: $sessionLogs")
-    return SingleSessionLog(sessionLogs.single())
+    return getSessionLogs(events).map { SingleSessionLog(it) }
   }
 
   class SingleSessionLog internal constructor(private val logEvent: LogEvent) {