IDE fatal error during tests generation #68

Author: Vassiliy-Kudryashov

utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/generator/TestGenerator.kt

@@ -23,6 +23,8 @@ import com.intellij.codeInsight.FileModificationService
 import com.intellij.ide.fileTemplates.FileTemplateManager
 import com.intellij.ide.fileTemplates.FileTemplateUtil
 import com.intellij.ide.fileTemplates.JavaTemplateUtil
+import com.intellij.openapi.application.runReadAction
+import com.intellij.openapi.application.runWriteAction
 import com.intellij.openapi.command.WriteCommandAction.runWriteCommandAction
 import com.intellij.openapi.command.executeCommand
 import com.intellij.openapi.components.service
@@ -44,39 +46,62 @@ import com.intellij.psi.codeStyle.JavaCodeStyleManager
 import com.intellij.psi.search.GlobalSearchScopesCore
 import com.intellij.testIntegration.TestIntegrationUtils
 import com.intellij.util.IncorrectOperationException
+import com.intellij.util.concurrency.AppExecutorUtil
 import com.intellij.util.io.exists
 import com.siyeh.ig.psiutils.ImportUtils
+import java.nio.file.Path
 import java.nio.file.Paths
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
 import org.jetbrains.kotlin.asJava.classes.KtUltraLightClass
 import org.jetbrains.kotlin.idea.core.ShortenReferences
 import org.jetbrains.kotlin.idea.core.getPackage
 import org.jetbrains.kotlin.idea.core.util.toPsiDirectory
 import org.jetbrains.kotlin.idea.util.ImportInsertHelperImpl
+import org.jetbrains.kotlin.idea.util.application.invokeLater
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtClass
 import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.psi.KtPsiFactory
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.scripting.resolve.classId
+import org.utbot.framework.plugin.api.util.UtContext
+import org.utbot.framework.plugin.api.util.withUtContext
 import org.utbot.intellij.plugin.error.showErrorDialogLater
 import org.utbot.intellij.plugin.ui.GenerateTestsModel
 import org.utbot.intellij.plugin.ui.SarifReportNotifier
 import org.utbot.intellij.plugin.ui.TestReportUrlOpeningListener
 import org.utbot.intellij.plugin.ui.TestsReportNotifier
 import org.utbot.intellij.plugin.ui.packageName
+import org.utbot.intellij.plugin.generator.TestGenerator.Target.THREAD_POOL
+import org.utbot.intellij.plugin.generator.TestGenerator.Target.EDT_LATER
+import org.utbot.intellij.plugin.generator.TestGenerator.Target.READ_ACTION
+import org.utbot.intellij.plugin.generator.TestGenerator.Target.WRITE_ACTION
 
 object TestGenerator {
-    fun generateTests(model: GenerateTestsModel, testCases: Map<PsiClass, List<UtTestCase>>) {
-        runWriteCommandAction(model.project, "Generate tests with UtBot", null, {
-            generateTestsInternal(model, testCases)
-        })
+    private enum class Target {THREAD_POOL, READ_ACTION, WRITE_ACTION, EDT_LATER}
+
+    private fun run(target: Target, runnable: Runnable) {
+        UtContext.currentContext()?.let {
+            when (target) {
+                THREAD_POOL -> AppExecutorUtil.getAppExecutorService().submit{
+                    withUtContext(it) {
+                        runnable.run()
+                    }
+                }
+                READ_ACTION -> runReadAction { withUtContext(it) { runnable.run() } }
+                WRITE_ACTION -> runWriteAction { withUtContext(it) { runnable.run() } }
+                EDT_LATER -> invokeLater { withUtContext(it) { runnable.run() } }
+            }
+        } ?: error("No context in thread ${Thread.currentThread()}")
     }
 
-    private fun generateTestsInternal(model: GenerateTestsModel, testCasesByClass: Map<PsiClass, List<UtTestCase>>) {
+    fun generateTests(model: GenerateTestsModel, testCasesByClass: Map<PsiClass, List<UtTestCase>>) {
         val baseTestDirectory = model.testSourceRoot?.toPsiDirectory(model.project)
             ?: return
         val allTestPackages = getPackageDirectories(baseTestDirectory)
+        val latch = CountDownLatch(testCasesByClass.size)
 
         for (srcClass in testCasesByClass.keys) {
             val testCases = testCasesByClass[srcClass] ?: continue
@@ -86,18 +111,37 @@ object TestGenerator {
                 val testDirectory = allTestPackages[classPackageName] ?: baseTestDirectory
                 val testClass = createTestClass(srcClass, testDirectory, model) ?: continue
                 val file = testClass.containingFile
-
-                addTestMethodsAndSaveReports(testClass, file, testCases, model)
+                runWriteCommandAction(model.project, "Generate tests with UtBot", null, {
+                    try {
+                        addTestMethodsAndSaveReports(testClass, file, testCases, model, latch)
+                    } catch (e: IncorrectOperationException) {
+                        showCreatingClassError(model.project, createTestClassName(srcClass))
+                    }
+                })
             } catch (e: IncorrectOperationException) {
                 showCreatingClassError(model.project, createTestClassName(srcClass))
             }
         }
+        run(READ_ACTION) {
+            val sarifReportsPath = model.testModule.getOrCreateSarifReportsPath(model.testSourceRoot)
+            run(THREAD_POOL) {
+                waitForCountDown(latch, model, sarifReportsPath)
+            }
+        }
+    }
 
-        mergeSarifReports(model)
+    private fun waitForCountDown(latch: CountDownLatch, model: GenerateTestsModel, sarifReportsPath : Path) {
+        try {
+            if (!latch.await(5, TimeUnit.SECONDS)) {
+                run(THREAD_POOL) { waitForCountDown(latch, model, sarifReportsPath) }
+            } else {
+                mergeSarifReports(model, sarifReportsPath)
+            }
+        } catch (ignored: InterruptedException) {
+        }
     }
 
-    private fun mergeSarifReports(model: GenerateTestsModel) {
-        val sarifReportsPath = model.testModule.getOrCreateSarifReportsPath(model.testSourceRoot)
+    private fun mergeSarifReports(model: GenerateTestsModel, sarifReportsPath : Path) {
         val sarifReports = sarifReportsPath.toFile()
             .walkTopDown()
             .filter { it.extension == "sarif" }
@@ -183,6 +227,7 @@ object TestGenerator {
         file: PsiFile,
         testCases: List<UtTestCase>,
         model: GenerateTestsModel,
+        reportsCountDown: CountDownLatch,
     ) {
         val selectedMethods = TestIntegrationUtils.extractClassMethods(testClass, false)
         val testFramework = model.testFramework
@@ -213,32 +258,48 @@ object TestGenerator {
         when (generator) {
             is ModelBasedTestCodeGenerator -> {
                 val editor = CodeInsightUtil.positionCursorAtLBrace(testClass.project, file, testClass)
-                val testsCodeWithTestReport = generator.generateAsStringWithTestReport(testCases)
-                val generatedTestsCode = testsCodeWithTestReport.generatedCode
+                //TODO Use PsiDocumentManager.getInstance(model.project).getDocument(file)
+                // if we don't want to open _all_ new files with tests in editor one-by-one
+                run(THREAD_POOL) {
+                    val testsCodeWithTestReport = generator.generateAsStringWithTestReport(testCases)
+                    val generatedTestsCode = testsCodeWithTestReport.generatedCode
+                    run(EDT_LATER) {
+                        run(WRITE_ACTION) {
+                            unblockDocument(testClass.project, editor.document)
+                            // TODO: JIRA:1246 - display warnings if we rewrite the file
+                            executeCommand(testClass.project, "Insert Generated Tests") {
+                                editor.document.setText(generatedTestsCode)
+                            }
+                            unblockDocument(testClass.project, editor.document)
+
+                            // after committing the document the `testClass` is invalid in PsiTree,
+                            // so we have to reload it from the corresponding `file`
+                            val testClassUpdated = (file as PsiClassOwner).classes.first() // only one class in the file
+
+                            // reformatting before creating reports due to
+                            // SarifReport requires the final version of the generated tests code
+                            runWriteCommandAction(testClassUpdated.project, "UtBot tests reformatting", null, {
+                                reformat(model, file, testClassUpdated)
+                            })
+                            unblockDocument(testClassUpdated.project, editor.document)
+
+                            // uploading formatted code
+                            val testsCodeWithTestReportFormatted =
+                                testsCodeWithTestReport.copy(generatedCode = file.text)
+
+                            // creating and saving reports
+                            saveSarifAndTestReports(
+                                testClassUpdated,
+                                testCases,
+                                model,
+                                testsCodeWithTestReportFormatted,
+                                reportsCountDown
+                            )
 
-                unblockDocument(testClass.project, editor.document)
-                // TODO: JIRA:1246 - display warnings if we rewrite the file
-                executeCommand(testClass.project, "Insert Generated Tests") {
-                    editor.document.setText(generatedTestsCode)
+                            unblockDocument(testClassUpdated.project, editor.document)
+                        }
+                    }
                 }
-                unblockDocument(testClass.project, editor.document)
-
-                // after committing the document the `testClass` is invalid in PsiTree,
-                // so we have to reload it from the corresponding `file`
-                val testClassUpdated = (file as PsiClassOwner).classes.first() // only one class in the file
-
-                // reformatting before creating reports due to
-                // SarifReport requires the final version of the generated tests code
-                reformat(model, file, testClassUpdated)
-                unblockDocument(testClassUpdated.project, editor.document)
-
-                // uploading formatted code
-                val testsCodeWithTestReportFormatted = testsCodeWithTestReport.copy(generatedCode = file.text)
-
-                // creating and saving reports
-                saveSarifAndTestReports(testClassUpdated, testCases, model, testsCodeWithTestReportFormatted)
-
-                unblockDocument(testClassUpdated.project, editor.document)
             }
             else -> TODO("Only model based code generator supported, but got: ${generator::class}")
         }
@@ -265,7 +326,8 @@ object TestGenerator {
         testClass: PsiClass,
         testCases: List<UtTestCase>,
         model: GenerateTestsModel,
-        testsCodeWithTestReport: TestsCodeWithTestReport
+        testsCodeWithTestReport: TestsCodeWithTestReport,
+        reportsCountDown: CountDownLatch
     ) {
         val project = model.project
         val generatedTestsCode = testsCodeWithTestReport.generatedCode
@@ -282,6 +344,8 @@ object TestGenerator {
                 message = "Cannot save Sarif report via generated tests: error occurred '${e.message}'",
                 title = "Failed to save Sarif report"
             )
+        } finally {
+            reportsCountDown.countDown()
         }
 
         try {