Collect missed models in stateAfter

Author: EgorkaKulikov

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

@@ -44,18 +44,13 @@ class SpringTestClassModelBuilder(val context: CgContext): TestClassModelBuilder
             for ((index, execution) in testSet.executions.withIndex()) {
                 execution.stateBefore.thisInstance?.let { model ->
                     thisInstances += model
-                    context.modelIds += model to model.modelId(index)
-
-                    val dependentModels = mutableSetOf<UtModel>()
-                    collectRecursively(model, dependentModels)
-
-                    dependentModels.forEach { model ->
-                        context.modelIds += model to model.modelId(index)
-                    }
-                    thisInstancesDependentModels += dependentModels
+                    thisInstancesDependentModels += collectByThisInstanceModel(model, index)
                 }
 
-                execution.stateAfter.thisInstance?.let { thisInstances += it }
+                execution.stateAfter.thisInstance?.let { model ->
+                    thisInstances += model
+                    thisInstancesDependentModels += collectByThisInstanceModel(model, index)
+                }
             }
         }
 
@@ -65,6 +60,19 @@ class SpringTestClassModelBuilder(val context: CgContext): TestClassModelBuilder
         return thisInstances.groupByClassId() to dependentMockModels.groupByClassId()
     }
 
+    private fun collectByThisInstanceModel(model: UtModel, executionIndex: Int): Set<UtModel> {
+        context.modelIds += model to model.modelId(executionIndex)
+
+        val dependentModels = mutableSetOf<UtModel>()
+        collectRecursively(model, dependentModels)
+
+        dependentModels.forEach { model ->
+            context.modelIds += model to model.modelId(executionIndex)
+        }
+
+        return dependentModels
+    }
+
     private fun Set<UtModel>.groupByClassId(): Map<ClassId, Set<UtModel>> {
         return this.groupBy { it.classId }.mapValues { it.value.toSet() }
     }