From 6a8504556dffc2b37ad56afc58d1f9273b04a997 Mon Sep 17 00:00:00 2001
From: Tomasz Godzik <tgodzik@virtuslab.com>
Date: Thu, 21 Sep 2023 15:14:31 +0200
Subject: [PATCH] bugfix: Try to fix Mill flaky test

Looks like Mill will sometimes throw an exception which is forwarded to Metals and causes tests to hange, this should make the tests stop hanging and make retries work.
---
 .../scala/tests/mill/MillServerCodeLensSuite.scala  | 10 +++++-----
 tests/unit/src/main/scala/tests/TestingServer.scala | 13 ++++++++++++-
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/tests/slow/src/test/scala/tests/mill/MillServerCodeLensSuite.scala b/tests/slow/src/test/scala/tests/mill/MillServerCodeLensSuite.scala
index 8003f47fbb3..6d25bf1d721 100644
--- a/tests/slow/src/test/scala/tests/mill/MillServerCodeLensSuite.scala
+++ b/tests/slow/src/test/scala/tests/mill/MillServerCodeLensSuite.scala
@@ -1,7 +1,5 @@
 package tests.mill
 
-import scala.concurrent.duration.Duration
-
 import scala.meta.internal.metals.ServerCommands
 import scala.meta.internal.metals.{BuildInfo => V}
 
@@ -12,8 +10,10 @@ import tests.MillServerInitializer
 class MillServerCodeLensSuite
     extends BaseCodeLensLspSuite("mill-server-lenses", MillServerInitializer) {
 
-  override def munitTimeout: Duration = Duration("4min")
-
+  /*
+   * There is some flakiness involved, possibly https://github.com/com-lihaoyi/mill/issues/2826
+   * We added some timeouts to make sure the test is correctly retried when fetching lenses.
+   */
   test("run-mill-lens", maxRetry = 3) {
     cleanWorkspace()
     writeLayout(
@@ -30,7 +30,7 @@ class MillServerCodeLensSuite
            |// no test lense as debug is not supported
            |class Foo extends munit.FunSuite {}
            |""".stripMargin,
-        V.scala213,
+        V.scala3,
         V.millVersion,
         includeMunit = true,
       )
diff --git a/tests/unit/src/main/scala/tests/TestingServer.scala b/tests/unit/src/main/scala/tests/TestingServer.scala
index a4ae0d0ade5..a86d2dca4a1 100644
--- a/tests/unit/src/main/scala/tests/TestingServer.scala
+++ b/tests/unit/src/main/scala/tests/TestingServer.scala
@@ -1034,6 +1034,7 @@ final case class TestingServer(
       filename: String,
       maxRetries: Int = 4,
   ): Future[List[l.CodeLens]] = {
+    Debug.printEnclosing(filename)
     val path = toPath(filename)
     val uri = path.toURI.toString
     val params = new CodeLensParams(new TextDocumentIdentifier(uri))
@@ -1048,9 +1049,18 @@ final case class TestingServer(
     var retries = maxRetries
     val codeLenses = Promise[List[l.CodeLens]]()
     val handler = { refreshCount: Int =>
+      scribe.info(s"Refreshing model for $filename")
       if (refreshCount > 0)
         for {
-          lenses <- fullServer.codeLens(params).asScala.map(_.asScala)
+          lenses <- fullServer
+            .codeLens(params)
+            .asScala
+            .map(_.asScala)
+            .withTimeout(10, util.concurrent.TimeUnit.SECONDS)
+            .recover { _ =>
+              scribe.info(s"Timeout for fetching lenses reached for $filename")
+              Nil
+            }
         } {
           if (lenses.nonEmpty) codeLenses.trySuccess(lenses.toList)
           else if (retries > 0) {
@@ -1072,6 +1082,7 @@ final case class TestingServer(
       // first compilation, to trigger the handler
       _ <- server.compilations.compileFile(path)
       lenses <- codeLenses.future
+        .withTimeout(60, util.concurrent.TimeUnit.SECONDS)
     } yield lenses
   }