Test framework installation before test generation #1570

utbot-cli-python/src/main/kotlin/org/utbot/cli/language/python/PythonGenerateTestsCommand.kt

@@ -244,6 +244,12 @@ class PythonGenerateTestsCommand : CliktCommand(
             checkingRequirementsAction = {
                 logger.info("Checking requirements...")
             },
+            installingRequirementsAction = {
+                logger.info("Installing requirements...")
+            },
+            testFrameworkInstallationAction = {
+                logger.info("Test framework installation...")
+            },
             requirementsAreNotInstalledAction = ::processMissingRequirements,
             startedLoadingPythonTypesAction = {
                 logger.info("Loading information about Python types...")

utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/PythonDialogProcessor.kt

@@ -51,20 +51,40 @@ object PythonDialogProcessor {
         focusedMethod: PyFunction?,
         file: PyFile
     ) {
-        val dialog = createDialog(project, functionsToShow, containingClass, focusedMethod, file)
-        if (!dialog.showAndGet()) {
-            return
+        val pythonPath = getPythonPath(functionsToShow)
+        if (pythonPath == null) {
+            showErrorDialogLater(
+                project,
+                message = "Couldn't find Python interpreter",
+                title = "Python test generation error"
+            )
+        } else {
+            val dialog = createDialog(
+                project,
+                functionsToShow,
+                containingClass,
+                focusedMethod,
+                file,
+                pythonPath,
+            )
+            if (!dialog.showAndGet()) {
+                return
+            }
+            createTests(project, dialog.model)
         }
+    }
 
-        createTests(project, dialog.model)
+    private fun getPythonPath(functionsToShow: Set<PyFunction>): String? {
+        return findSrcModule(functionsToShow).sdk?.homePath
     }
 
     private fun createDialog(
         project: Project,
         functionsToShow: Set<PyFunction>,
         containingClass: PyClass?,
         focusedMethod: PyFunction?,
-        file: PyFile
+        file: PyFile,
+        pythonPath: String,
     ): PythonDialogWindow {
         val srcModule = findSrcModule(functionsToShow)
         val testModules = srcModule.testModules(project)
@@ -85,6 +105,7 @@ object PythonDialogProcessor {
                 DEFAULT_TIMEOUT_FOR_RUN_IN_MILLIS,
                 visitOnlySpecifiedSource = false,
                 cgLanguageAssistant = PythonCgLanguageAssistant,
+                pythonPath = pythonPath,
             )
         )
     }
@@ -119,15 +140,7 @@ object PythonDialogProcessor {
                     return
                 }
                 try {
-                    val pythonPath = model.srcModule.sdk?.homePath
-                    if (pythonPath == null) {
-                        showErrorDialogLater(
-                            project,
-                            message = "Couldn't find Python interpreter",
-                            title = "Python test generation error"
-                        )
-                        return
-                    }
+
                     val methods = findSelectedPythonMethods(model)
                     if (methods == null) {
                         showErrorDialogLater(
@@ -138,7 +151,7 @@ object PythonDialogProcessor {
                         return
                     }
                     processTestGeneration(
-                        pythonPath = pythonPath,
+                        pythonPath = model.pythonPath,
                         pythonFilePath = model.file.virtualFile.path,
                         pythonFileContent = getContentFromPyFile(model.file),
                         directoriesForSysPath = model.directoriesForSysPath,
@@ -151,8 +164,10 @@ object PythonDialogProcessor {
                         visitOnlySpecifiedSource = model.visitOnlySpecifiedSource,
                         isCanceled = { indicator.isCanceled },
                         checkingRequirementsAction = { indicator.text = "Checking requirements" },
+                        installingRequirementsAction = { indicator.text = "Installing requirements..." },
+                        testFrameworkInstallationAction = { indicator.text = "Test framework installation" },
                         requirementsAreNotInstalledAction = {
-                            askAndInstallRequirementsLater(model.project, pythonPath)
+                            askAndInstallRequirementsLater(model.project, model.pythonPath)
                             PythonTestGenerationProcessor.MissingRequirementsActionResult.NOT_INSTALLED
                         },
                         startedLoadingPythonTypesAction = { indicator.text = "Loading information about Python types" },

utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/PythonDialogWindow.kt

@@ -5,8 +5,10 @@ import com.intellij.openapi.ui.ComboBox
 import com.intellij.openapi.ui.DialogPanel
 import com.intellij.openapi.ui.DialogWrapper
 import com.intellij.openapi.ui.ValidationInfo
+import com.intellij.ui.ColoredListCellRenderer
 import com.intellij.ui.ContextHelpLabel
 import com.intellij.ui.JBIntSpinner
+import com.intellij.ui.SimpleTextAttributes
 import com.intellij.ui.components.Panel
 import com.intellij.ui.layout.CellBuilder
 import com.intellij.ui.layout.Row
@@ -17,16 +19,15 @@ import com.jetbrains.python.refactoring.classes.PyMemberInfoStorage
 import com.jetbrains.python.refactoring.classes.membersManager.PyMemberInfo
 import com.jetbrains.python.refactoring.classes.ui.PyMemberSelectionTable
 import org.utbot.framework.UtSettings
-import org.utbot.framework.plugin.api.CodeGenerationSettingItem
+import org.utbot.framework.codegen.domain.TestFramework
 import java.awt.BorderLayout
 import java.util.concurrent.TimeUnit
-import javax.swing.DefaultComboBoxModel
-import javax.swing.JCheckBox
-import javax.swing.JComponent
-import javax.swing.JPanel
 import org.utbot.intellij.plugin.ui.components.TestSourceDirectoryChooser
+import org.utbot.intellij.plugin.ui.utils.createTestFrameworksRenderer
+import javax.swing.*
 
 
+private const val WILL_BE_INSTALLED_LABEL = " (will be installed)"
 private const val MINIMUM_TIMEOUT_VALUE_IN_SECONDS = 1
 
 class PythonDialogWindow(val model: PythonTestsModel) : DialogWrapper(model.project) {
@@ -55,12 +56,16 @@ class PythonDialogWindow(val model: PythonTestsModel) : DialogWrapper(model.proj
     private lateinit var panel: DialogPanel
 
     init {
-        title = "Generate Tests With UtBot"
+        title = "Generate Tests with UnitTestBot"
         isResizable = false
+
+        model.cgLanguageAssistant.getLanguageTestFrameworkManager().testFrameworks.forEach {
+            it.isInstalled = it.isInstalled || checkModuleIsInstalled(model.pythonPath, it.mainPackage)
+        }
+
         init()
     }
 
-    @Suppress("UNCHECKED_CAST")
     override fun createCenterPanel(): JComponent {
 
         panel = panel {
@@ -69,7 +74,7 @@ class PythonDialogWindow(val model: PythonTestsModel) : DialogWrapper(model.proj
             }
             row("Test framework:") {
                 makePanelWithHelpTooltip(
-                    testFrameworks as ComboBox<CodeGenerationSettingItem>,
+                    testFrameworks,
                     null
                 )
             }
@@ -100,9 +105,14 @@ class PythonDialogWindow(val model: PythonTestsModel) : DialogWrapper(model.proj
         }
 
         updateFunctionsTable()
+        updateTestFrameworksList()
         return panel
     }
 
+    private fun updateTestFrameworksList() {
+        testFrameworks.renderer = createTestFrameworksRenderer(WILL_BE_INSTALLED_LABEL)
+    }
+
     private fun globalPyFunctionsToPyMemberInfo(
         project: Project,
         functions: Collection<PyFunction>

utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/PythonTestsModel.kt

@@ -23,6 +23,7 @@ class PythonTestsModel(
     var timeoutForRun: Long,
     var visitOnlySpecifiedSource: Boolean,
     val cgLanguageAssistant: CgLanguageAssistant,
+    val pythonPath: String,
 ) : BaseTestsModel(
     project,
     srcModule,

utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/Utils.kt

@@ -4,6 +4,7 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.ProjectFileIndex
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiElement
+import org.utbot.python.utils.RequirementsUtils
 import kotlin.random.Random
 
 inline fun <reified T : PsiElement> getContainingElement(
@@ -28,3 +29,7 @@ fun generateRandomString(length: Int): String {
         .map { Random.nextInt(0, charPool.size).let { charPool[it] } }
         .joinToString("")
 }
+
+fun checkModuleIsInstalled(pythonPath: String, moduleName: String): Boolean {
+    return RequirementsUtils.requirementsAreInstalled(pythonPath, listOf(moduleName))
+}

utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/ui/GenerateTestsDialogWindow.kt

@@ -116,7 +116,6 @@ import org.utbot.framework.plugin.api.utils.MOCKITO_EXTENSIONS_FILE_CONTENT
 import org.utbot.framework.plugin.api.utils.MOCKITO_EXTENSIONS_FOLDER
 import org.utbot.framework.plugin.api.utils.MOCKITO_MOCKMAKER_FILE_NAME
 import org.utbot.framework.util.Conflict
-import org.utbot.intellij.plugin.generator.UtTestsDialogProcessor
 import org.utbot.intellij.plugin.models.GenerateTestsModel
 import org.utbot.intellij.plugin.models.id
 import org.utbot.intellij.plugin.models.jUnit4LibraryDescriptor
@@ -132,6 +131,7 @@ import org.utbot.intellij.plugin.ui.components.TestFolderComboWithBrowseButton
 import org.utbot.intellij.plugin.ui.utils.LibrarySearchScope
 import org.utbot.intellij.plugin.ui.utils.addSourceRootIfAbsent
 import org.utbot.intellij.plugin.ui.utils.allLibraries
+import org.utbot.intellij.plugin.ui.utils.createTestFrameworksRenderer
 import org.utbot.intellij.plugin.ui.utils.findFrameworkLibrary
 import org.utbot.intellij.plugin.ui.utils.findParametrizedTestsLibrary
 import org.utbot.intellij.plugin.ui.utils.getOrCreateTestResourcesPath
@@ -985,17 +985,7 @@ class GenerateTestsDialogWindow(val model: GenerateTestsModel) : DialogWrapper(m
 
         testFrameworks.model = DefaultComboBoxModel(enabledTestFrameworks.toTypedArray())
         testFrameworks.item = if (currentFrameworkItem in enabledTestFrameworks) currentFrameworkItem else defaultItem
-        testFrameworks.renderer = object : ColoredListCellRenderer<TestFramework>() {
-            override fun customizeCellRenderer(
-                list: JList<out TestFramework>, value: TestFramework,
-                index: Int, selected: Boolean, hasFocus: Boolean
-            ) {
-                this.append(value.displayName, SimpleTextAttributes.REGULAR_ATTRIBUTES)
-                if (!value.isInstalled) {
-                    this.append(WILL_BE_INSTALLED_LABEL, SimpleTextAttributes.ERROR_ATTRIBUTES)
-                }
-            }
-        }
+        testFrameworks.renderer = createTestFrameworksRenderer(WILL_BE_INSTALLED_LABEL)
 
         currentFrameworkItem = testFrameworks.item
     }

utbot-python/src/main/kotlin/org/utbot/python/PythonTestGenerationProcessor.kt

@@ -22,8 +22,8 @@ import org.utbot.python.typing.MypyAnnotations
 import org.utbot.python.typing.PythonTypesStorage
 import org.utbot.python.typing.StubFileFinder
 import org.utbot.python.utils.Cleaner
+import org.utbot.python.utils.RequirementsUtils.installRequirements
 import org.utbot.python.utils.RequirementsUtils.requirementsAreInstalled
-import org.utbot.python.utils.TemporaryFileManager
 import org.utbot.python.utils.getLineOfFunction
 import java.io.File
 import java.nio.file.Path
@@ -49,6 +49,8 @@ object PythonTestGenerationProcessor {
         withMinimization: Boolean = true,
         isCanceled: () -> Boolean = { false },
         checkingRequirementsAction: () -> Unit = {},
+        installingRequirementsAction: () -> Unit = {},
+        testFrameworkInstallationAction: () -> Unit = {},
         requirementsAreNotInstalledAction: () -> MissingRequirementsActionResult = {
             MissingRequirementsActionResult.NOT_INSTALLED
         },
@@ -63,11 +65,14 @@ object PythonTestGenerationProcessor {
         Cleaner.restart()
 
         try {
-            TemporaryFileManager.setup()
-
+            if (!testFramework.isInstalled) {
+                testFrameworkInstallationAction()
+                installRequirements(pythonPath, listOf(testFramework.mainPackage))
+            }
             if (!doNotCheckRequirements) {
                 checkingRequirementsAction()
                 if (!requirementsAreInstalled(pythonPath)) {
+                    installingRequirementsAction()
                     val result = requirementsAreNotInstalledAction()
                     if (result == MissingRequirementsActionResult.NOT_INSTALLED)
                         return

utbot-python/src/main/kotlin/org/utbot/python/framework/codegen/model/PythonDomain.kt

@@ -48,6 +48,10 @@ object Pytest : TestFramework(displayName = "pytest", id = "pytest") {
 }
 
 object Unittest : TestFramework(displayName = "Unittest", id = "Unittest") {
+    init {
+        isInstalled = true
+    }
+
     override val testSuperClass: ClassId = PythonClassId("unittest.TestCase")
     override val mainPackage: String = "unittest"
     override val assertionsClass: ClassId = PythonClassId("self")

utbot-python/src/main/kotlin/org/utbot/python/utils/RequirementsUtils.kt

@@ -14,25 +14,33 @@ object RequirementsUtils {
             ?: error("Didn't find /check_requirements.py")
 
     fun requirementsAreInstalled(pythonPath: String): Boolean {
+        return requirementsAreInstalled(pythonPath, requirements)
+    }
+
+    fun requirementsAreInstalled(pythonPath: String, requirementList: List<String>): Boolean {
         val requirementsScript =
             TemporaryFileManager.createTemporaryFile(requirementsScriptContent, tag = "requirements")
         val result = runCommand(
             listOf(
                 pythonPath,
                 requirementsScript.path
-            ) + requirements
+            ) + requirementList
         )
         return result.exitValue == 0
     }
 
     fun installRequirements(pythonPath: String): CmdResult {
+        return installRequirements(pythonPath, requirements)
+    }
+
+    fun installRequirements(pythonPath: String, moduleNames: List<String>): CmdResult {
         return runCommand(
             listOf(
                 pythonPath,
                 "-m",
                 "pip",
                 "install"
-            ) + requirements
+            ) + moduleNames
         )
     }
 }

utbot-python/src/main/kotlin/org/utbot/python/utils/TemporaryFileManager.kt

@@ -6,10 +6,10 @@ import java.nio.file.Path
 import java.nio.file.Paths
 
 object TemporaryFileManager {
-    private lateinit var tmpDirectory: Path
+    private var tmpDirectory: Path
     private var nextId = 0
 
-    fun setup() {
+    init {
         tmpDirectory = FileUtil.createTempDirectory("python-test-generation-${nextId++}")
         Cleaner.addFunction { tmpDirectory.toFile().deleteRecursively() }
     }

utbot-ui-commons/src/main/kotlin/org/utbot/intellij/plugin/ui/utils/DialogWindowUtils.kt

@@ -0,0 +1,20 @@
+package org.utbot.intellij.plugin.ui.utils
+
+import com.intellij.ui.ColoredListCellRenderer
+import com.intellij.ui.SimpleTextAttributes
+import org.utbot.framework.codegen.domain.TestFramework
+import javax.swing.JList
+
+fun createTestFrameworksRenderer(willBeInstalledLabel: String): ColoredListCellRenderer<TestFramework> {
+    return object : ColoredListCellRenderer<TestFramework>() {
+        override fun customizeCellRenderer(
+            list: JList<out TestFramework>, value: TestFramework,
+            index: Int, selected: Boolean, hasFocus: Boolean
+        ) {
+            this.append(value.displayName, SimpleTextAttributes.REGULAR_ATTRIBUTES)
+            if (!value.isInstalled) {
+                this.append(willBeInstalledLabel, SimpleTextAttributes.ERROR_ATTRIBUTES)
+            }
+        }
+    }
+}