From 633ccca4c201281ab1675d2f407be5ec8165b36e Mon Sep 17 00:00:00 2001
From: Tomasz Godzik <tgodzik@virtuslab.com>
Date: Thu, 10 Oct 2024 16:06:01 +0200
Subject: [PATCH] improvement: Retry code actions in Scala CLI suite

---
 .../scalacli/BaseScalaCLIActionSuite.scala    |  2 ++
 .../tests/scalacli/ScalaCliActionsSuite.scala |  4 +++
 .../codeactions/BaseCodeActionLspSuite.scala  | 34 ++++++++++++-------
 3 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/tests/slow/src/main/scala/tests/scalacli/BaseScalaCLIActionSuite.scala b/tests/slow/src/main/scala/tests/scalacli/BaseScalaCLIActionSuite.scala
index 0a393090de0..b5caa0d2485 100644
--- a/tests/slow/src/main/scala/tests/scalacli/BaseScalaCLIActionSuite.scala
+++ b/tests/slow/src/main/scala/tests/scalacli/BaseScalaCLIActionSuite.scala
@@ -66,6 +66,7 @@ class BaseScalaCLIActionSuite(name: String)
       changeFile: String => String = identity,
       expectError: Boolean = false,
       filterAction: CodeAction => Boolean = _ => true,
+      retryAction: Int = 0,
   )(implicit loc: Location): Unit = {
 
     val path = toPath(fileName)
@@ -96,6 +97,7 @@ class BaseScalaCLIActionSuite(name: String)
       expectError,
       filterAction,
       overrideLayout = layout,
+      retryAction,
     )
   }
 }
diff --git a/tests/slow/src/test/scala/tests/scalacli/ScalaCliActionsSuite.scala b/tests/slow/src/test/scala/tests/scalacli/ScalaCliActionsSuite.scala
index c6e306bd8a7..3926de64f68 100644
--- a/tests/slow/src/test/scala/tests/scalacli/ScalaCliActionsSuite.scala
+++ b/tests/slow/src/test/scala/tests/scalacli/ScalaCliActionsSuite.scala
@@ -53,6 +53,8 @@ class ScalaCliActionsSuite
     scalaCliOptions = List("--actions", "-S", scalaVersion),
     expectNoDiagnostics = false,
     selectedActionIndex = 1,
+    // Scala CLI doesn't publish everything with the normal diagnostics, but later
+    retryAction = 5,
   )
 
   checkScalaCLI(
@@ -83,6 +85,8 @@ class ScalaCliActionsSuite
     scalaCliOptions = List("--actions", "-S", scalaVersion),
     expectNoDiagnostics = false,
     selectedActionIndex = 1,
+    // Scala CLI doesn't publish everything with the normal diagnostics, but later
+    retryAction = 5,
   )
 
   checkNoActionScalaCLI(
diff --git a/tests/unit/src/main/scala/tests/codeactions/BaseCodeActionLspSuite.scala b/tests/unit/src/main/scala/tests/codeactions/BaseCodeActionLspSuite.scala
index 6350030fa41..078e4c0c384 100644
--- a/tests/unit/src/main/scala/tests/codeactions/BaseCodeActionLspSuite.scala
+++ b/tests/unit/src/main/scala/tests/codeactions/BaseCodeActionLspSuite.scala
@@ -59,6 +59,7 @@ abstract class BaseCodeActionLspSuite(
       expectError: Boolean = false,
       filterAction: CodeAction => Boolean = _ => true,
       overrideLayout: Option[String] = None,
+      retryAction: Int = 0,
   )(implicit loc: Location): Unit = {
     val scalacOptionsJson =
       if (scalacOptions.nonEmpty)
@@ -88,6 +89,7 @@ abstract class BaseCodeActionLspSuite(
       changeFile,
       expectError,
       filterAction,
+      retryAction,
     )
   }
 
@@ -105,6 +107,7 @@ abstract class BaseCodeActionLspSuite(
       changeFile: String => String = identity,
       expectError: Boolean = false,
       filterAction: CodeAction => Boolean = _ => true,
+      retryAction: Int = 0,
   )(implicit loc: Location): Unit = {
     val files = FileLayout.mapFromString(layout)
     val (path, input) = files
@@ -120,6 +123,24 @@ abstract class BaseCodeActionLspSuite(
       if (renamePath.nonEmpty) input.replace("<<", "").replace(">>", "")
       else expectedCode
 
+    def assertActionsWithRetry(
+        retry: Int = retryAction
+    ): Future[List[CodeAction]] = {
+      server
+        .assertCodeAction(
+          path,
+          changeFile(input),
+          expectedActions,
+          kind,
+          filterAction = filterAction,
+        )
+        .recoverWith {
+          case _: Throwable if retry > 0 =>
+            Thread.sleep(1000)
+            assertActionsWithRetry(retry - 1)
+          case _: Throwable if expectError => Future.successful(Nil)
+        }
+    }
     test(name) {
       cleanWorkspace()
       for {
@@ -135,18 +156,7 @@ abstract class BaseCodeActionLspSuite(
           path,
           changeFile(input).replace("<<", "").replace(">>", ""),
         )
-        codeActions <-
-          server
-            .assertCodeAction(
-              path,
-              changeFile(input),
-              expectedActions,
-              kind,
-              filterAction = filterAction,
-            )
-            .recover {
-              case _: Throwable if expectError => Nil
-            }
+        codeActions <- assertActionsWithRetry()
         _ <- client.applyCodeAction(selectedActionIndex, codeActions, server)
         _ <- server.didSave(newPath) { _ =>
           if (newPath != path)