Fix code generator pipeline for cases when util class is generated separately

Author: ArsenHD

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

@@ -65,19 +65,19 @@ class CodeGenerator(
     fun generateAsStringWithTestReport(
         testSets: Collection<UtMethodTestSet>,
         testClassCustomName: String? = null,
-    ): CodeGenerationResult {
+    ): CodeGeneratorResult {
         val cgTestSets = testSets.map { CgMethodTestSet(it) }.toList()
         return generateAsStringWithTestReport(cgTestSets, testClassCustomName)
     }
 
     private fun generateAsStringWithTestReport(
         cgTestSets: List<CgMethodTestSet>,
         testClassCustomName: String? = null,
-    ): CodeGenerationResult = withCustomContext(testClassCustomName) {
+    ): CodeGeneratorResult = withCustomContext(testClassCustomName) {
         context.withTestClassFileScope {
             val testClassModel = TestClassModel.fromTestSets(classUnderTest, cgTestSets)
             val testClassFile = CgTestClassConstructor(context).construct(testClassModel)
-            CodeGenerationResult(
+            CodeGeneratorResult(
                 generatedCode = renderClassFile(testClassFile),
                 utilClassKind = UtilClassKind.fromCgContextOrNull(context),
                 testsGenerationReport = testClassFile.testsGenerationReport,
@@ -117,7 +117,7 @@ class CodeGenerator(
  * @property testsGenerationReport some info about test generation process
  * @property mockFrameworkUsed flag indicating whether any mock objects have been created during code generation ot not
  */
-data class CodeGenerationResult(
+data class CodeGeneratorResult(
     val generatedCode: String,
     // null if no util class needed, e.g. when we are generating utils directly into test class
     val utilClassKind: UtilClassKind?,

utbot-framework/src/test/kotlin/org/utbot/framework/codegen/CompilationAndRunUtils.kt

@@ -14,6 +14,13 @@ data class ClassUnderTest(
 
 private val logger = KotlinLogging.logger {}
 
+fun writeFile(fileContents: String, targetFile: File): File {
+    val targetDir = targetFile.parentFile
+    targetDir.mkdirs()
+    targetFile.writeText(fileContents)
+    return targetFile
+}
+
 fun writeTest(
     testContents: String,
     testClassName: String,
@@ -26,13 +33,10 @@ fun writeTest(
         File(buildDirectory.toFile(), "${testClassName.substringAfterLast(".")}${generatedLanguage.extension}")
     )
 
-    val targetDir = classUnderTest.generatedTestFile.parentFile
-    targetDir.mkdirs()
     logger.info {
         "File size for ${classUnderTest.testClassSimpleName}: ${FileUtil.byteCountToDisplaySize(testContents.length.toLong())}"
     }
-    classUnderTest.generatedTestFile.writeText(testContents)
-    return classUnderTest.generatedTestFile
+    return writeFile(testContents, classUnderTest.generatedTestFile)
 }
 
 fun compileTests(

utbot-framework/src/test/kotlin/org/utbot/framework/codegen/TestCodeGeneratorPipeline.kt

@@ -9,6 +9,9 @@ import org.utbot.common.packageName
 import org.utbot.examples.TestFrameworkConfiguration
 import org.utbot.framework.codegen.ExecutionStatus.SUCCESS
 import org.utbot.framework.codegen.model.CodeGenerator
+import org.utbot.framework.codegen.model.CodeGeneratorResult
+import org.utbot.framework.codegen.model.UtilClassKind
+import org.utbot.framework.codegen.model.UtilClassKind.Companion.UT_UTILS_CLASS_NAME
 import org.utbot.framework.plugin.api.CodegenLanguage
 import org.utbot.framework.plugin.api.ExecutableId
 import org.utbot.framework.plugin.api.MockFramework
@@ -18,6 +21,8 @@ import org.utbot.framework.plugin.api.util.UtContext
 import org.utbot.framework.plugin.api.util.description
 import org.utbot.framework.plugin.api.util.id
 import org.utbot.framework.plugin.api.util.withUtContext
+import java.io.File
+import java.nio.file.Path
 import kotlin.reflect.KClass
 
 private val logger = KotlinLogging.logger {}
@@ -74,7 +79,8 @@ class TestCodeGeneratorPipeline(private val testFrameworkConfiguration: TestFram
             val codegenLanguage = testFrameworkConfiguration.codegenLanguage
             val parametrizedTestSource = testFrameworkConfiguration.parametrizedTestSource
 
-            val testClass = callToCodeGenerator(testSets, classUnderTest)
+            val codeGenerationResult = callToCodeGenerator(testSets, classUnderTest)
+            val testClass = codeGenerationResult.generatedCode
 
             // actual number of the tests in the generated testClass
             val generatedMethodsCount = testClass
@@ -148,17 +154,33 @@ class TestCodeGeneratorPipeline(private val testFrameworkConfiguration: TestFram
 
                 val testClassName = classPipeline.retrieveTestClassName("BrokenGeneratedTest")
                 val generatedTestFile = writeTest(testClass, testClassName, buildDirectory, codegenLanguage)
+                val generatedUtilClassFile = codeGenerationResult.utilClassKind?.writeUtilClassToFile(buildDirectory, codegenLanguage)
 
                 logger.error("Broken test has been written to the file: [$generatedTestFile]")
+                if (generatedUtilClassFile != null) {
+                    logger.error("Util class for the broken test has been written to the file: [$generatedUtilClassFile]")
+                }
                 logger.error("Failed configuration: $testFrameworkConfiguration")
 
                 throw it
             }
 
-            classPipeline.stageContext = copy(data = testClass, stages = stages + information.completeStage())
+            classPipeline.stageContext = copy(data = codeGenerationResult, stages = stages + information.completeStage())
         }
     }
 
+    private fun UtilClassKind.writeUtilClassToFile(buildDirectory: Path, language: CodegenLanguage): File {
+        val utilClassFile = File(buildDirectory.toFile(), "$UT_UTILS_CLASS_NAME${language.extension}")
+        val utilClassText = getUtilClassText(language)
+        return writeFile(utilClassText, utilClassFile)
+    }
+
+    private data class GeneratedTestClassInfo(
+        val testClassName: String,
+        val generatedTestFile: File,
+        val generatedUtilClassFile: File?
+    )
+
     @Suppress("UNCHECKED_CAST")
     private fun processCompilationStages(classesPipelines: List<ClassPipeline>) {
         val information = StageExecutionInformation(Compilation)
@@ -168,24 +190,34 @@ class TestCodeGeneratorPipeline(private val testFrameworkConfiguration: TestFram
         val codegenLanguage = testFrameworkConfiguration.codegenLanguage
 
         val testClassesNamesToTestGeneratedTests = classesPipelines.map { classPipeline ->
-            val testClass = classPipeline.stageContext.data as String
+            val codeGeneratorResult = classPipeline.stageContext.data as CodeGeneratorResult//String
+            val testClass = codeGeneratorResult.generatedCode
+
             val testClassName = classPipeline.retrieveTestClassName("GeneratedTest")
             val generatedTestFile = writeTest(testClass, testClassName, buildDirectory, codegenLanguage)
+            val generatedUtilClassFile = codeGeneratorResult.utilClassKind?.writeUtilClassToFile(buildDirectory, codegenLanguage)
 
             logger.info("Test has been written to the file: [$generatedTestFile]")
+            if (generatedUtilClassFile != null) {
+                logger.info("Util class for the test has been written to the file: [$generatedUtilClassFile]")
+            }
 
-            testClassName to generatedTestFile
+            GeneratedTestClassInfo(testClassName, generatedTestFile, generatedUtilClassFile)
         }
 
+        val sourceFiles = mutableListOf<String>().apply {
+            this += testClassesNamesToTestGeneratedTests.map { it.generatedTestFile.absolutePath }
+            this += testClassesNamesToTestGeneratedTests.mapNotNull { it.generatedUtilClassFile?.absolutePath }
+        }
         compileTests(
             "$buildDirectory",
-            testClassesNamesToTestGeneratedTests.map { it.second.absolutePath },
+            sourceFiles,
             codegenLanguage
         )
 
-        testClassesNamesToTestGeneratedTests.zip(classesPipelines) { testClassNameToTest, classPipeline ->
+        testClassesNamesToTestGeneratedTests.zip(classesPipelines) { generatedTestClassInfo, classPipeline ->
             classPipeline.stageContext = classPipeline.stageContext.copy(
-                data = CompilationResult("$buildDirectory", testClassNameToTest.first),
+                data = CompilationResult("$buildDirectory", generatedTestClassInfo.testClassName),
                 stages = classPipeline.stageContext.stages + information.completeStage()
             )
         }
@@ -223,7 +255,7 @@ class TestCodeGeneratorPipeline(private val testFrameworkConfiguration: TestFram
     private fun callToCodeGenerator(
         testSets: List<UtMethodTestSet>,
         classUnderTest: KClass<*>
-    ): String {
+    ): CodeGeneratorResult {
         val params = mutableMapOf<ExecutableId, List<String>>()
 
         val codeGenerator = with(testFrameworkConfiguration) {
@@ -243,7 +275,7 @@ class TestCodeGeneratorPipeline(private val testFrameworkConfiguration: TestFram
         }
         val testClassCustomName = "${classUnderTest.java.simpleName}GeneratedTest"
 
-        return codeGenerator.generateAsString(testSets, testClassCustomName)
+        return codeGenerator.generateAsStringWithTestReport(testSets, testClassCustomName)
     }
 
     private fun checkPipelinesResults(classesPipelines: List<ClassPipeline>) {

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

@@ -57,7 +57,7 @@ import org.utbot.framework.codegen.ParametrizedTestSource
 import org.utbot.framework.codegen.RegularImport
 import org.utbot.framework.codegen.StaticImport
 import org.utbot.framework.codegen.model.CodeGenerator
-import org.utbot.framework.codegen.model.CodeGenerationResult
+import org.utbot.framework.codegen.model.CodeGeneratorResult
 import org.utbot.framework.codegen.model.UtilClassKind
 import org.utbot.framework.codegen.model.UtilClassKind.Companion.UT_UTILS_CLASS_NAME
 import org.utbot.framework.codegen.model.constructor.tree.TestsGenerationReport
@@ -96,7 +96,7 @@ object CodeGenerationController {
         var requiredUtilClassKind: UtilClassKind? = null
         var mockFrameworkUsed: Boolean = false
 
-        fun onTestClassGenerated(result: CodeGenerationResult) {
+        fun onTestClassGenerated(result: CodeGeneratorResult) {
             requiredUtilClassKind = maxOfNullable(requiredUtilClassKind, result.utilClassKind)
             mockFrameworkUsed = maxOf(mockFrameworkUsed, result.mockFrameworkUsed)
         }
@@ -615,7 +615,7 @@ object CodeGenerationController {
         testClass: PsiClass,
         testSets: List<UtMethodTestSet>,
         model: GenerateTestsModel,
-        testsCodeWithTestReport: CodeGenerationResult,
+        testsCodeWithTestReport: CodeGeneratorResult,
     ) {
         val project = model.project
         val generatedTestsCode = testsCodeWithTestReport.generatedCode