[Headless mode] Revive headless mode

build.gradle.kts

@@ -2,7 +2,6 @@ import org.jetbrains.changelog.markdownToHTML
 import org.jetbrains.intellij.platform.gradle.IntelliJPlatformType
 import org.jetbrains.intellij.platform.gradle.TestFrameworkType
 import org.jetbrains.intellij.platform.gradle.models.ProductRelease
-import org.jetbrains.intellij.platform.gradle.tasks.RunIdeTask
 import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
 import java.io.FileOutputStream
 import java.net.URL
@@ -46,6 +45,8 @@ repositories {
     // See https://plugins.jetbrains.com/docs/intellij/tools-intellij-platform-gradle-plugin-repositories-extension.html#default-repositories
     intellijPlatform {
         defaultRepositories()
+        intellijDependencies()
+        localPlatformArtifacts()
     }
     maven("https://packages.jetbrains.team/maven/p/ij/intellij-dependencies")
     maven("https://www.jetbrains.com/intellij-repository/snapshots")
@@ -65,21 +66,23 @@ repositories {
     }
 }
 
+var hasGrazieAccess: SourceSet? = null
+
 if (spaceCredentialsProvided()) {
     // Add the new source set
-    val hasGrazieAccess = sourceSets.create("hasGrazieAccess")
+    hasGrazieAccess = sourceSets.create("hasGrazieAccess")
     // add output of main source set to new source set class path
-    hasGrazieAccess.compileClasspath += sourceSets.main.get().output
+    hasGrazieAccess!!.compileClasspath += sourceSets.main.get().output
 
     // register feature variant
-    java.registerFeature(hasGrazieAccess.name) {
-        usingSourceSet(hasGrazieAccess)
+    java.registerFeature(hasGrazieAccess!!.name) {
+        usingSourceSet(hasGrazieAccess!!)
     }
 
     // Add the dependencies for the new source set
     dependencies {
-        add(hasGrazieAccess.implementationConfigurationName, "org.jetbrains.kotlinx:kotlinx-coroutines-core:1.7.3")
-        add(hasGrazieAccess.implementationConfigurationName, "org.jetbrains.research:grazie-test-generation:$grazieTestGenerationVersion")
+        add(hasGrazieAccess!!.implementationConfigurationName, "org.jetbrains.kotlinx:kotlinx-coroutines-core:1.7.3")
+        add(hasGrazieAccess!!.implementationConfigurationName, "org.jetbrains.research:grazie-test-generation:$grazieTestGenerationVersion")
     }
 
     tasks.register("checkCredentials") {
@@ -89,24 +92,25 @@ if (spaceCredentialsProvided()) {
             ).files()
     }
 
-    tasks.named(hasGrazieAccess.jarTaskName).configure {
+    tasks.named(hasGrazieAccess!!.jarTaskName).configure {
         dependsOn("checkCredentials")
     }
 
-    tasks.named<Jar>(hasGrazieAccess.jarTaskName) {
+    tasks.named<Jar>(hasGrazieAccess!!.jarTaskName) {
         exclude("**/plugin.xml")
     }
 
     // add build of new source set as the part of UI testing
     tasks.prepareTestSandbox.configure {
-        dependsOn(hasGrazieAccess.jarTaskName)
+        dependsOn(hasGrazieAccess!!.jarTaskName)
         from(
             tasks
-                .getByName(hasGrazieAccess.jarTaskName)
+                .getByName(hasGrazieAccess!!.jarTaskName)
                 .outputs.files.asPath,
         ) { into("TestSpark/lib") }
 
-        hasGrazieAccess.runtimeClasspath
+        hasGrazieAccess!!
+            .runtimeClasspath
             .elements
             .get()
             .forEach {
@@ -115,14 +119,15 @@ if (spaceCredentialsProvided()) {
     }
     // add build of new source set as the part of pluginBuild process
     tasks.prepareSandbox.configure {
-        dependsOn(hasGrazieAccess.jarTaskName)
+        dependsOn(hasGrazieAccess!!.jarTaskName)
         from(
             tasks
-                .getByName(hasGrazieAccess.jarTaskName)
+                .getByName(hasGrazieAccess!!.jarTaskName)
                 .outputs.files.asPath,
         ) { into("TestSpark/lib") }
 
-        hasGrazieAccess.runtimeClasspath
+        hasGrazieAccess!!
+            .runtimeClasspath
             .elements
             .get()
             .forEach {
@@ -464,7 +469,30 @@ fun String?.orDefault(default: String): String = this ?: default
  * @param out The output directory for the project.
  * @param enableCoverage flag to enable/disable coverage computation
  */
-tasks.create<RunIdeTask>("headless") {
+val headless by intellijPlatformTesting.runIde.registering {
+    if (hasGrazieAccess == null) {
+        println("ERROR: Grazie access not found, skipping headless test generation")
+        return@registering
+    }
+    prepareSandboxTask.configure {
+        dependsOn(hasGrazieAccess!!.jarTaskName)
+        from(
+            tasks
+                .getByName(hasGrazieAccess!!.jarTaskName)
+                .outputs.files.asPath,
+        ) { into("TestSpark/lib") }
+
+        hasGrazieAccess!!
+            .runtimeClasspath
+            .elements
+            .get()
+            .forEach {
+                from(it.asFile.absolutePath) { into("TestSpark/lib") }
+            }
+    }
+    type = IntelliJPlatformType.IntellijIdeaCommunity
+    version = properties("platformVersion")
+
     val root: String? by project
     val file: String? by project
     val cut: String? by project
@@ -476,15 +504,21 @@ tasks.create<RunIdeTask>("headless") {
     val out: String? by project
     val enableCoverage: String? by project
 
-    args = listOfNotNull("testspark", root, file, cut, cp, junitv, llm, token, prompt, out, enableCoverage.orDefault("false"))
+    task {
+        dependsOn("buildPlugin")
+        jvmArgumentProviders +=
+            CommandLineArgumentProvider {
+                listOf(
+                    "-Xmx16G",
+                    "-Djava.awt.headless=true",
+                    "--add-exports",
+                    "java.base/jdk.internal.vm=ALL-UNNAMED",
+                    "-Didea.system.path",
+                )
+            }
 
-    jvmArgs(
-        "-Xmx16G",
-        "-Djava.awt.headless=true",
-        "--add-exports",
-        "java.base/jdk.internal.vm=ALL-UNNAMED",
-        "-Didea.system.path",
-    )
+        args = listOfNotNull("testspark", root, file, cut, cp, junitv, llm, token, prompt, out, enableCoverage.orDefault("false"))
+    }
 }
 
 fun spaceCredentialsProvided() = spaceUsername.isNotEmpty() && spacePassword.isNotEmpty()

gradle/wrapper/gradle-wrapper.properties

@@ -1,5 +1,5 @@
 distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
-distributionUrl=https\://services.gradle.org/distributions/gradle-8.8-bin.zip
+distributionUrl=https\://services.gradle.org/distributions/gradle-8.9-bin.zip
 zipStoreBase=GRADLE_USER_HOME
 zipStorePath=wrapper/dists

src/main/kotlin/org/jetbrains/research/testspark/appstarter/ProjectApplicationUtils.kt

@@ -13,6 +13,8 @@ import com.intellij.openapi.application.ModalityState
 import com.intellij.openapi.application.ex.ApplicationManagerEx
 import com.intellij.openapi.progress.util.ProgressIndicatorBase
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.project.ProjectManager
+import com.intellij.openapi.startup.StartupManager
 import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.util.io.FileUtil
 import com.intellij.openapi.vfs.LocalFileSystem
@@ -25,6 +27,7 @@ import org.jetbrains.plugins.gradle.GradleCommandLineProjectConfigurator
 import org.slf4j.LoggerFactory
 import java.nio.file.Path
 import java.util.function.Predicate
+import kotlin.coroutines.resume
 import kotlin.coroutines.suspendCoroutine
 
 class ProjectConfiguratorException : Exception {
@@ -170,16 +173,18 @@ object ProjectApplicationUtils {
         logger.info("Closing project $project...")
         ApplicationManager.getApplication().assertIsNonDispatchThread()
         // ToDo: move headless mode to another branch
-//        ApplicationManager.getApplication().invokeAndWait {
-//            ProjectManagerEx.getInstanceEx().forceCloseProject(project)
-//        }
+        ApplicationManager.getApplication().invokeAndWait {
+            ProjectManager.getInstance().closeAndDispose(project)
+            // TODO: what is ProjectManagerEx?
+            // ProjectManagerEx.getInstanceEx().forceCloseProject(project)
+        }
     }
 
     private suspend fun waitAllStartupActivitiesPassed(project: Project): Unit =
         suspendCoroutine {
             logger.info("Waiting all startup activities passed $project...")
             // ToDo: move headless mode to another branch
-//        StartupManager.getInstance(project).runAfterOpened { it.resume(Unit) }
+            StartupManager.getInstance(project).runAfterOpened { it.resume(Unit) }
             waitForInvokeLaterActivities()
         }
 

src/main/kotlin/org/jetbrains/research/testspark/appstarter/TestSparkStarter.kt

@@ -14,10 +14,12 @@ import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiJavaFile
 import com.intellij.psi.PsiManager
 import kotlinx.serialization.ExperimentalSerializationApi
+import org.jetbrains.research.testspark.actions.controllers.IndicatorController
 import org.jetbrains.research.testspark.bundles.llm.LLMDefaultsBundle
 import org.jetbrains.research.testspark.core.data.JUnitVersion
 import org.jetbrains.research.testspark.core.data.TestGenerationData
 import org.jetbrains.research.testspark.core.monitor.DefaultErrorMonitor
+import org.jetbrains.research.testspark.core.test.SupportedLanguage
 import org.jetbrains.research.testspark.core.test.TestCompiler
 import org.jetbrains.research.testspark.core.test.data.CodeType
 import org.jetbrains.research.testspark.data.FragmentToTestData
@@ -31,7 +33,8 @@ import org.jetbrains.research.testspark.tools.TestProcessor
 import org.jetbrains.research.testspark.tools.TestsExecutionResultManager
 import org.jetbrains.research.testspark.tools.ToolUtils
 import org.jetbrains.research.testspark.tools.factories.TestCompilerFactory
-import org.jetbrains.research.testspark.tools.llm.Llm
+import org.jetbrains.research.testspark.tools.llm.generation.LLMProcessManager
+import org.jetbrains.research.testspark.tools.llm.generation.PromptManager
 import java.io.File
 import java.nio.file.Path
 import java.nio.file.Paths
@@ -73,7 +76,7 @@ class TestSparkStarter : ApplicationStarter {
         val runCoverage = args[10].toBoolean()
 
         val testsExecutionResultManager = TestsExecutionResultManager()
-
+        val indicatorController = IndicatorController()
         println("Test generation requested for $projectPath")
 
         // remove the `.idea` folder in the $projectPath if exists
@@ -179,14 +182,13 @@ class TestSparkStarter : ApplicationStarter {
                         }
                         // Create a process Manager
                         val llmProcessManager =
-                            Llm()
-                                .getLLMProcessManager(
-                                    project,
-                                    psiHelper!!,
-                                    targetPsiClass.textRange.startOffset,
-                                    testSamplesCode = "", // we don't provide samples to LLM
-                                    projectSDKPath = projectSDKPath,
-                                )
+                            LLMProcessManager(
+                                project,
+                                language = SupportedLanguage.Java,
+                                PromptManager(project, psiHelper!!, targetPsiClass.textRange.startOffset),
+                                testSamplesCode = "", // we don't provide samples to LLM
+                                projectSDKPath = projectSDKPath,
+                            )
 
                         println("[TestSpark Starter] Starting the test generation process")
                         // Start test generation
@@ -206,6 +208,7 @@ class TestSparkStarter : ApplicationStarter {
                                 packageName,
                                 projectContext,
                                 testGenerationData,
+                                indicatorController,
                                 errorMonitor,
                                 testsExecutionResultManager,
                             )