Generate separated parametrized tests for test sets with different result type

Author: sofurihafe

utbot-framework-api/src/main/kotlin/org/utbot/framework/plugin/api/Api.kt

@@ -29,6 +29,7 @@ import org.utbot.framework.plugin.api.util.isPrimitive
 import org.utbot.framework.plugin.api.util.jClass
 import org.utbot.framework.plugin.api.util.longClassId
 import org.utbot.framework.plugin.api.util.method
+import org.utbot.framework.plugin.api.util.objectClassId
 import org.utbot.framework.plugin.api.util.primitiveTypeJvmNameOrNull
 import org.utbot.framework.plugin.api.util.shortClassId
 import org.utbot.framework.plugin.api.util.toReferenceTypeBytecodeSignature
@@ -100,20 +101,63 @@ data class UtMethodTestSet(
 
 data class CgMethodTestSet private constructor(
     val executableId: ExecutableId,
-    val executions: List<UtExecution> = emptyList(),
     val jimpleBody: JimpleBody? = null,
     val errors: Map<String, Int> = emptyMap(),
-    val clustersInfo: List<Pair<UtClusterInfo?, IntRange>> = listOf(null to executions.indices)
+    val clustersInfo: List<Pair<UtClusterInfo?, IntRange>>,
 ) {
+    var executions: List<UtExecution> = emptyList()
+        private set
+
     constructor(from: UtMethodTestSet) : this(
         from.method.callable.executableId,
-        from.executions,
         from.jimpleBody,
         from.errors,
         from.clustersInfo
-    )
+    ) {
+        executions = from.executions
+    }
+
+    /**
+     * Splits [CgMethodTestSet] into separate test sets having
+     * unique result model [ClassId] in each subset.
+     */
+    fun splitExecutionsByResult() : List<CgMethodTestSet> {
+        val successfulExecutions = executions.filter { it.result is UtExecutionSuccess }
+        val executionsByResult: Map<ClassId, List<UtExecution>> =
+            if (successfulExecutions.any()) {
+                successfulExecutions.groupBy { (it.result as UtExecutionSuccess).model.classId }
+            } else {
+                mapOf(objectClassId to executions)
+            }
+
+        return executionsByResult.map{ (_, executions) -> substituteExecutions(executions) }
+    }
+
+    /**
+     * Finds a [ClassId] of all result models in executions.
+     *
+     * Tries to find an unique result type in testSets or
+     * gets executable return type.
+     */
+    fun resultType(): ClassId {
+        val successfulExecutions = executions.filter { it.result is UtExecutionSuccess }
+        return if (successfulExecutions.any()) {
+            successfulExecutions
+                .map { (it.result as UtExecutionSuccess).model.classId }
+                .distinct()
+                .singleOrNull()
+                ?: executableId.returnType
+        } else {
+            executableId.returnType
+        }
+    }
+
+    private fun substituteExecutions(newExecutions: List<UtExecution>): CgMethodTestSet =
+        this.copy().apply { executions = newExecutions }
+
 }
 
+
 data class Step(
     val stmt: Stmt,
     val depth: Int,

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/model/constructor/tree/CgMethodConstructor.kt

@@ -1265,9 +1265,7 @@ internal class CgMethodConstructor(val context: CgContext) : CgContextOwner by c
                 currentMethodParameters[CgParameterKind.Argument(index)] = argument.parameter
             }
 
-            val method = currentExecutable!!
-            val expectedResultClassId = wrapTypeIfRequired(method.returnType)
-
+            val expectedResultClassId = wrapTypeIfRequired(testSet.resultType())
             if (expectedResultClassId != voidClassId) {
                 val wrappedType = wrapIfPrimitive(expectedResultClassId)
                 //We are required to wrap the type of expected result if it is primitive

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/model/constructor/tree/CgTestClassConstructor.kt

@@ -83,7 +83,7 @@ internal class CgTestClassConstructor(val context: CgContext) :
             return null
         }
 
-        val (methodUnderTest, executions, _, _, clustersInfo) = testSet
+        val (methodUnderTest, _, _, clustersInfo) = testSet
         val regions = mutableListOf<CgRegion<CgMethod>>()
         val requiredFields = mutableListOf<CgParameterDeclaration>()
 
@@ -94,7 +94,7 @@ internal class CgTestClassConstructor(val context: CgContext) :
                     emptyLineIfNeeded()
                     for (i in executionIndices) {
                         runCatching {
-                            currentTestCaseTestMethods += methodConstructor.createTestMethod(methodUnderTest, executions[i])
+                            currentTestCaseTestMethods += methodConstructor.createTestMethod(methodUnderTest, testSet.executions[i])
                         }.onFailure { e -> processFailure(testSet, e) }
                     }
                     val clusterHeader = clusterSummary?.header
@@ -107,22 +107,9 @@ internal class CgTestClassConstructor(val context: CgContext) :
                 }
             }
             ParametrizedTestSource.PARAMETRIZE -> {
-                runCatching {
-                    val dataProviderMethodName = nameGenerator.dataProviderMethodNameFor(testSet.executableId)
-
-                    val parameterizedTestMethod =
-                        methodConstructor.createParameterizedTestMethod(testSet, dataProviderMethodName)
-
-                    requiredFields += parameterizedTestMethod.requiredFields
-
-                    cgDataProviderMethods +=
-                        methodConstructor.createParameterizedTestDataProvider(testSet, dataProviderMethodName)
-
-                    regions += CgSimpleRegion(
-                        "Parameterized test for method ${methodUnderTest.displayName}",
-                        listOf(parameterizedTestMethod),
-                    )
-                }.onFailure { error -> processFailure(testSet, error) }
+                for (currentTestSet in testSet.splitExecutionsByResult()) {
+                    createParametrizedTestAndDataProvider(currentTestSet, requiredFields, regions, methodUnderTest)
+                }
             }
         }
 
@@ -141,6 +128,30 @@ internal class CgTestClassConstructor(val context: CgContext) :
             .merge(failure.description, 1, Int::plus)
     }
 
+    private fun createParametrizedTestAndDataProvider(
+        testSet: CgMethodTestSet,
+        requiredFields: MutableList<CgParameterDeclaration>,
+        regions: MutableList<CgRegion<CgMethod>>,
+        methodUnderTest: ExecutableId,
+    ) {
+        runCatching {
+            val dataProviderMethodName = nameGenerator.dataProviderMethodNameFor(testSet.executableId)
+
+            val parameterizedTestMethod =
+                methodConstructor.createParameterizedTestMethod(testSet, dataProviderMethodName)
+
+            requiredFields += parameterizedTestMethod.requiredFields
+
+            cgDataProviderMethods +=
+                methodConstructor.createParameterizedTestDataProvider(testSet, dataProviderMethodName)
+
+            regions += CgSimpleRegion(
+                "Parameterized test for method ${methodUnderTest.displayName}",
+                listOf(parameterizedTestMethod),
+            )
+        }.onFailure { error -> processFailure(testSet, error) }
+    }
+
     // TODO: collect imports of util methods
     private fun createUtilMethods(): List<CgUtilMethod> {
         val utilMethods = mutableListOf<CgUtilMethod>()

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

@@ -129,10 +129,12 @@ class TestCodeGeneratorPipeline(private val testFrameworkConfiguration: TestFram
                     "Errors regions has been generated: $errorText"
                 }
 
-                require(generatedMethodsCount == expectedNumberOfGeneratedMethods) {
-                    "Something went wrong during the code generation for ${classUnderTest.simpleName}. " +
-                            "Expected to generate $expectedNumberOfGeneratedMethods test methods, " +
-                            "but got only $generatedMethodsCount"
+                if (parametrizedTestSource != ParametrizedTestSource.PARAMETRIZE) {
+                    require(generatedMethodsCount == expectedNumberOfGeneratedMethods) {
+                        "Something went wrong during the code generation for ${classUnderTest.simpleName}. " +
+                                "Expected to generate $expectedNumberOfGeneratedMethods test methods, " +
+                                "but got only $generatedMethodsCount"
+                    }
                 }
             }.onFailure {
                 val classes = listOf(classPipeline).retrieveClasses()