IJPL-578 preload kotlin script libraries in advance in kotlin highlighting tests to avoid flaky failures

If the libraries are not loaded, they will be added to the project right from the highlighting pass. This will trigger WA, which cancels current highlighting and restarts it at some point in the future. Test may check the state before the highlighting is restarted (this will be the test failure), ot after it is restarted (test pass).

This commit eliminates the highlighting restart

GitOrigin-RevId: a79d699dca94f5eac2690160ca87c2f576636c28

plugins/kotlin/base/scripting/src/org/jetbrains/kotlin/idea/core/script/ScriptConfigurationManager.kt

@@ -3,13 +3,16 @@
 package org.jetbrains.kotlin.idea.core.script
 
 import com.intellij.ide.scratch.ScratchUtil
+import com.intellij.lang.injection.InjectedLanguageManager
 import com.intellij.openapi.components.service
 import com.intellij.openapi.components.serviceIfCreated
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.Sdk
 import com.intellij.openapi.vfs.StandardFileSystems
 import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
+import com.intellij.psi.PsiRecursiveElementVisitor
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.util.io.URLUtil
 import org.jetbrains.annotations.TestOnly
@@ -129,10 +132,28 @@ interface ScriptConfigurationManager {
 
         fun toVfsRoots(roots: Iterable<File>): List<VirtualFile> = roots.mapNotNull { classpathEntryToVfs(it.toPath()) }
 
+        @Suppress("TestOnlyProblems")
         @TestOnly
         fun updateScriptDependenciesSynchronously(file: PsiFile) {
             // TODO: review the usages of this method
-            defaultScriptingSupport(file.project).updateScriptDependenciesSynchronously(file)
+            val defaultScriptingSupport = defaultScriptingSupport(file.project)
+            when(file) {
+                is KtFile -> {
+                    defaultScriptingSupport.updateScriptDependenciesSynchronously(file)
+                }
+                else -> {
+                    val project = file.project
+                    val injectedLanguageManager = InjectedLanguageManager.getInstance(project)
+                    object : PsiRecursiveElementVisitor() {
+                        override fun visitElement(element: PsiElement) {
+                            injectedLanguageManager.enumerate(element) { psi, _ ->
+                                defaultScriptingSupport.updateScriptDependenciesSynchronously(psi)
+                            }
+                          super.visitElement(element)
+                        }
+                    }.visitFile(file)
+                }
+            }
         }
 
         private fun defaultScriptingSupport(project: Project) =

plugins/kotlin/idea/tests/test/org/jetbrains/kotlin/idea/highlighter/AbstractHighlightingTest.java

@@ -13,7 +13,11 @@ import com.intellij.testFramework.fixtures.impl.CodeInsightTestFixtureImpl;
 import kotlin.Unit;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.kotlin.idea.base.test.InTextDirectivesUtils;
-import org.jetbrains.kotlin.idea.test.*;
+import org.jetbrains.kotlin.idea.core.script.ScriptConfigurationManager;
+import org.jetbrains.kotlin.idea.test.ConfigLibraryUtil;
+import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase;
+import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCaseKt;
+import org.jetbrains.kotlin.idea.test.TagsTestDataUtil;
 
 import java.io.File;
 import java.util.List;
@@ -27,6 +31,7 @@ public abstract class AbstractHighlightingTest extends KotlinLightCodeInsightFix
     public static final String NO_CHECK_WEAK_WARNINGS_PREFIX = "// NO_CHECK_WEAK_WARNINGS";
     public static final String NO_CHECK_WARNINGS_PREFIX = "// NO_CHECK_WARNINGS";
     public static final String EXPECTED_DUPLICATED_HIGHLIGHTING_PREFIX = "// EXPECTED_DUPLICATED_HIGHLIGHTING";
+    public static final String LOAD_SCRIPT_DEFINITIONS_DIRECTIVE = "// LOAD_SCRIPT_DEFINITIONS";
     public static final String ALLOW_DOC_CHANGE_PREFIX = "// ALLOW_DOC_CHANGE";
     public static final String TOOL_PREFIX = "// TOOL:";
 
@@ -38,6 +43,10 @@ public abstract class AbstractHighlightingTest extends KotlinLightCodeInsightFix
 
         KotlinLightCodeInsightFixtureTestCaseKt.withCustomCompilerOptions(fileText, getProject(), getModule(), () ->
         {
+            if (InTextDirectivesUtils.isDirectiveDefined(fileText, LOAD_SCRIPT_DEFINITIONS_DIRECTIVE)) {
+                ScriptConfigurationManager.Companion.updateScriptDependenciesSynchronously(myFixture.getFile());
+            }
+
             ExpectedHighlightingData data = new ExpectedHighlightingData(myFixture.getEditor().getDocument(), checkWarnings, checkWeakWarnings, checkInfos);
             if (checkInfos) data.checkSymbolNames();
             ((CodeInsightTestFixtureImpl) myFixture).canChangeDocumentDuringHighlighting(allowDocChange);

plugins/kotlin/idea/tests/testData/highlighter/KotlinInJavaInjection.java

@@ -1,4 +1,5 @@
 // ALLOW_DOC_CHANGE (project model is updated since the injected snippet is teated as a script)
+// LOAD_SCRIPT_DEFINITIONS
 public class KotlinInJavaInjection {
     public void inject() {
         String kotlin =