Result class improved

Author: EgorkaKulikov

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

@@ -122,20 +122,39 @@ data class CgMethodTestSet private constructor(
      * 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 (executions.any { it.result is UtExecutionSuccess })
-                executions
-                    .filter { it.result is UtExecutionSuccess }
-                    .groupBy { (it.result as UtExecutionSuccess).model.classId }
-            else {
+            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 }
+
 }
 
 

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