Some corrections

Author: EgorkaKulikov

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

@@ -45,7 +45,7 @@ class SpringTestClassModelBuilder(val context: CgContext): TestClassModelBuilder
     private fun collectModelsForClassVariables(testSets: List<CgMethodTestSet>): SpringSpecificInformation {
         val thisInstanceModels = mutableSetOf<UtModelWrapper>()
         val thisInstancesDependentModels = mutableSetOf<UtModelWrapper>()
-        val applicationContextModels = mutableSetOf<UtModelWrapper>()
+        val stateBeforeDependentModels = mutableSetOf<UtModelWrapper>()
 
         with(context) {
             for ((testSetIndex, testSet) in testSets.withIndex()) {
@@ -56,16 +56,13 @@ class SpringTestClassModelBuilder(val context: CgContext): TestClassModelBuilder
                                 .filterNotNull()
                                 .forEach { model ->
                                     thisInstanceModels += model.wrap()
-                                    thisInstancesDependentModels += collectByThisInstanceModel(model)
+                                    thisInstancesDependentModels += collectByModel(model)
 
                                 }
-                            //TODO: собрать рекурснивно все модели, как в моках, и из них уже выбрать что надо
-                            execution.stateBefore.parameters + setOf(execution.stateBefore.thisInstance)
-                                .filterNotNull()
-                                .forEach { model ->
-                                    applicationContextModels += collectApplicationContextModels(model)
-                                }
 
+                            (execution.stateBefore.parameters + setOf(execution.stateBefore.thisInstance))
+                                .filterNotNull()
+                                .forEach { model -> stateBeforeDependentModels += collectByModel(model) }
                         }
                     }
                 }
@@ -78,14 +75,18 @@ class SpringTestClassModelBuilder(val context: CgContext): TestClassModelBuilder
                     cgModel.model.isMockModel() && cgModel !in thisInstanceModels
                 }
 
+        val applicationContextModels = stateBeforeDependentModels
+                .filter { it.model is UtSpringContextModel }
+                .toSet()
+
         return SpringSpecificInformation(
             thisInstanceModels.groupByClassId(),
             dependentMockModels.groupByClassId(),
-            mapOf(applicationContextClassId to applicationContextModels),
+            applicationContextModels.groupByClassId(),
         )
     }
 
-    private fun collectByThisInstanceModel(model: UtModel): Set<UtModelWrapper> {
+    private fun collectByModel(model: UtModel): Set<UtModelWrapper> {
         val dependentModels = mutableSetOf<UtModelWrapper>()
         collectRecursively(model, dependentModels)
 
@@ -147,21 +148,4 @@ class SpringTestClassModelBuilder(val context: CgContext): TestClassModelBuilder
             //Python, JavaScript, Go models are not required in Spring
         }
     }
-
-    private fun collectApplicationContextModels(model: UtModel): Set<UtModelWrapper> {
-        if (model !is UtAssembleModel) {
-            return emptySet()
-        }
-
-        with (context) {
-            val innerModels = mutableSetOf<UtModel>()
-            model.instantiationCall.instance?.let { innerModels += it }
-            model.modificationsChain.mapNotNull { it.instance }.forEach { innerModels += it }
-
-            return innerModels
-                .filterIsInstance<UtSpringContextModel>()
-                .map { it.wrap() }
-                .toSet()
-        }
-    }
 }