Enlarge ModelId in codegen with testSetId (#2059)

Author: EgorkaKulikov

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/domain/Domain.kt

@@ -780,13 +780,16 @@ object SpringBoot : DependencyInjectionFramework(
 /**
  * Extended id of [UtModel], unique for whole test set.
  *
- * Allows to distinguish models from different executions,
- * even if they have the same value of `UtModel.id`.
+ * Allows distinguishing models from different executions and test sets,
+ * even if they have the same value of `UtModel.id` that is allowed.
  */
-data class ModelId(
+data class ModelId private constructor(
     private val id: Int?,
     private val executionId: Int,
-)
-
-fun UtModel.withExecutionId(executionId: Int = -1) = ModelId(this.idOrNull(), executionId)
+    private val testSetId: Int,
+) {
+    companion object {
+        fun create(model: UtModel, executionId: Int = -1, testSetId: Int = -1) = ModelId(model.idOrNull(), executionId, testSetId)
+    }
+}
 

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/domain/context/CgContext.kt

@@ -30,7 +30,6 @@ import org.utbot.framework.codegen.domain.builtin.UtilClassFileMethodProvider
 import org.utbot.framework.codegen.domain.builtin.UtilMethodProvider
 import org.utbot.framework.codegen.domain.models.SimpleTestClassModel
 import org.utbot.framework.codegen.domain.models.CgParameterKind
-import org.utbot.framework.codegen.domain.withExecutionId
 import org.utbot.framework.codegen.services.access.Block
 import org.utbot.framework.codegen.tree.EnvironmentFieldStateCache
 import org.utbot.framework.codegen.tree.importIfNeeded
@@ -571,7 +570,7 @@ data class CgContext(
         }
     }
 
-    override fun getIdByModel(model: UtModel): ModelId = modelIds.getOrPut(model) { model.withExecutionId() }
+    override fun getIdByModel(model: UtModel): ModelId = modelIds.getOrPut(model) { ModelId.create(model) }
 
     private fun createClassIdForNestedClass(testClassModel: SimpleTestClassModel): ClassId {
         val simpleName = "${testClassModel.classUnderTest.simpleName}Test"

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/domain/models/builders/SpringTestClassModelBuilder.kt

@@ -1,10 +1,10 @@
 package org.utbot.framework.codegen.domain.models.builders
 
+import org.utbot.framework.codegen.domain.ModelId
 import org.utbot.framework.codegen.domain.context.CgContext
 import org.utbot.framework.codegen.domain.models.CgMethodTestSet
 import org.utbot.framework.codegen.domain.models.ClassModels
 import org.utbot.framework.codegen.domain.models.SpringTestClassModel
-import org.utbot.framework.codegen.domain.withExecutionId
 import org.utbot.framework.plugin.api.ClassId
 import org.utbot.framework.plugin.api.UtArrayModel
 import org.utbot.framework.plugin.api.UtAssembleModel
@@ -39,17 +39,15 @@ class SpringTestClassModelBuilder(val context: CgContext): TestClassModelBuilder
         val thisInstances = mutableSetOf<UtModel>()
         val thisInstancesDependentModels = mutableSetOf<UtModel>()
 
-        for (testSet in testSets) {
-            for ((index, execution) in testSet.executions.withIndex()) {
-                execution.stateBefore.thisInstance?.let { model ->
-                    thisInstances += model
-                    thisInstancesDependentModels += collectByThisInstanceModel(model, index)
-                }
+        for ((testSetIndex, testSet) in testSets.withIndex()) {
+            for ((executionIndex, execution) in testSet.executions.withIndex()) {
 
-                execution.stateAfter.thisInstance?.let { model ->
-                    thisInstances += model
-                    thisInstancesDependentModels += collectByThisInstanceModel(model, index)
-                }
+                setOf(execution.stateBefore.thisInstance, execution.stateAfter.thisInstance)
+                    .filterNotNull()
+                    .forEach { model ->
+                        thisInstances += model
+                        thisInstancesDependentModels += collectByThisInstanceModel(model, executionIndex, testSetIndex)
+                    }
             }
         }
 
@@ -59,14 +57,14 @@ class SpringTestClassModelBuilder(val context: CgContext): TestClassModelBuilder
         return thisInstances.groupByClassId() to dependentMockModels.groupByClassId()
     }
 
-    private fun collectByThisInstanceModel(model: UtModel, executionIndex: Int): Set<UtModel> {
-        context.modelIds[model] = model.withExecutionId(executionIndex)
+    private fun collectByThisInstanceModel(model: UtModel, executionIndex: Int, testSetIndex: Int): Set<UtModel> {
+        context.modelIds[model] = ModelId.create(model, executionIndex, testSetIndex)
 
         val dependentModels = mutableSetOf<UtModel>()
         collectRecursively(model, dependentModels)
 
         dependentModels.forEach { model ->
-            context.modelIds[model] = model.withExecutionId(executionIndex)
+            context.modelIds[model] = ModelId.create(model, executionIndex, testSetIndex)
         }
 
         return dependentModels