Refactor annotated class variables processing

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

@@ -20,11 +20,14 @@ import org.utbot.framework.codegen.domain.builtin.UtilClassFileMethodProvider
 import org.utbot.framework.codegen.domain.builtin.UtilMethodProvider
 import org.utbot.framework.codegen.domain.models.*
 import org.utbot.framework.codegen.services.access.Block
+import org.utbot.framework.codegen.services.access.CgFieldStateManager
 import org.utbot.framework.codegen.tree.EnvironmentFieldStateCache
 import org.utbot.framework.codegen.tree.importIfNeeded
 import org.utbot.framework.plugin.api.BuiltinClassId
 import org.utbot.framework.plugin.api.ClassId
 import org.utbot.framework.codegen.services.language.CgLanguageAssistant
+import org.utbot.framework.codegen.tree.fieldmanager.CgAbstractClassFieldManager
+import org.utbot.framework.codegen.tree.fieldmanager.CgClassFieldManager
 import org.utbot.framework.plugin.api.CodegenLanguage
 import org.utbot.framework.plugin.api.ExecutableId
 import org.utbot.framework.plugin.api.FieldId
@@ -242,6 +245,12 @@ interface CgContextOwner {
      */
     var successfulExecutionsModels: List<UtModel>
 
+    /**
+     * Managers to process annotated fields of the class under test
+     * relevant for the current generation type.
+     */
+    val relevantFieldManagers: MutableList<CgAbstractClassFieldManager>
+
     fun block(init: () -> Unit): Block {
         val prevBlock = currentBlock
         return try {
@@ -508,6 +517,7 @@ class CgContext(
         RuntimeExceptionTestsBehaviour.defaultItem,
     override val hangingTestsTimeout: HangingTestsTimeout = HangingTestsTimeout(),
     override val enableTestsTimeout: Boolean = true,
+    override val relevantFieldManagers: MutableList<CgAbstractClassFieldManager> = mutableListOf(),
     override var containsReflectiveCall: Boolean = false,
 ) : CgContextOwner {
     override lateinit var statesCache: EnvironmentFieldStateCache

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

@@ -20,20 +20,3 @@ class SimpleTestClassModel(
     nestedClasses: List<SimpleTestClassModel> = listOf(),
 ): TestClassModel(classUnderTest, methodTestSets, nestedClasses)
 
-/**
- * Extended [SimpleTestClassModel] for Spring analysis reasons
- */
-class SpringTestClassModel(
-    classUnderTest: ClassId,
-    methodTestSets: List<CgMethodTestSet>,
-    nestedClasses: List<SimpleTestClassModel>,
-    val springSpecificInformation: SpringSpecificInformation,
-): TestClassModel(classUnderTest, methodTestSets, nestedClasses)
-
-class SpringSpecificInformation(
-    val thisInstanceModels: TypedModelWrappers,
-    val thisInstanceDependentMocks: TypedModelWrappers,
-    val thisInstanceDependentSpies: TypedModelWrappers,
-    val autowiredFromContextModels: TypedModelWrappers,
-    val entityManagerModels: TypedModelWrappers,
-)

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

@@ -1,12 +1,14 @@
 package org.utbot.framework.codegen.domain.models.builders
 
-import org.utbot.framework.codegen.domain.context.CgContext
+import org.utbot.framework.codegen.domain.UtModelWrapper
 import org.utbot.framework.codegen.domain.models.CgMethodTestSet
 import org.utbot.framework.codegen.domain.models.SimpleTestClassModel
 import org.utbot.framework.plugin.api.ClassId
 import org.utbot.framework.plugin.api.util.enclosingClass
 
-open class SimpleTestClassModelBuilder(context: CgContext): TestClassModelBuilder() {
+typealias TypedModelWrappers = Map<ClassId, Set<UtModelWrapper>>
+
+open class SimpleTestClassModelBuilder: TestClassModelBuilder() {
     override fun createTestClassModel(
         classUnderTest: ClassId,
         testSets: List<CgMethodTestSet>,

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/generator/CodeGenerator.kt

@@ -10,7 +10,7 @@ open class CodeGenerator(params: CodeGeneratorParams): AbstractCodeGenerator(par
     protected val classUnderTest: ClassId = params.classUnderTest
 
     override fun generate(testSets: List<CgMethodTestSet>): CodeGeneratorResult {
-        val testClassModel = SimpleTestClassModelBuilder(context).createTestClassModel(classUnderTest, testSets)
+        val testClassModel = SimpleTestClassModelBuilder().createTestClassModel(classUnderTest, testSets)
 
         logger.info { "Code generation phase started at ${now()}" }
         val astConstructor = CgSimpleTestClassConstructor(context)

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/generator/SpringCodeGenerator.kt

@@ -2,7 +2,7 @@ package org.utbot.framework.codegen.generator
 
 import org.utbot.framework.codegen.domain.context.CgContext
 import org.utbot.framework.codegen.domain.models.CgMethodTestSet
-import org.utbot.framework.codegen.domain.models.builders.SpringTestClassModelBuilder
+import org.utbot.framework.codegen.domain.models.builders.SimpleTestClassModelBuilder
 import org.utbot.framework.codegen.services.language.CgLanguageAssistant
 import org.utbot.framework.codegen.tree.CgCustomAssertConstructor
 import org.utbot.framework.codegen.tree.CgSpringIntegrationTestClassConstructor
@@ -39,7 +39,7 @@ class SpringCodeGenerator(
     private val classUnderTest: ClassId = params.classUnderTest
 
     override fun generate(testSets: List<CgMethodTestSet>): CodeGeneratorResult {
-        val testClassModel = SpringTestClassModelBuilder(context).createTestClassModel(classUnderTest, testSets)
+        val testClassModel = SimpleTestClassModelBuilder().createTestClassModel(classUnderTest, testSets)
 
         logger.info { "Code generation phase started at ${now()}" }
         val astConstructor = when (springTestType) {

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/tree/CgAbstractSpringTestClassConstructor.kt

@@ -2,9 +2,8 @@ package org.utbot.framework.codegen.tree
 
 import org.utbot.framework.codegen.domain.builtin.TestClassUtilMethodProvider
 import org.utbot.framework.codegen.domain.context.CgContext
-import org.utbot.framework.codegen.domain.models.AnnotationTarget.*
+import org.utbot.framework.codegen.domain.models.AnnotationTarget.Method
 import org.utbot.framework.codegen.domain.models.CgClassBody
-import org.utbot.framework.codegen.domain.models.CgDeclaration
 import org.utbot.framework.codegen.domain.models.CgFieldDeclaration
 import org.utbot.framework.codegen.domain.models.CgFrameworkUtilMethod
 import org.utbot.framework.codegen.domain.models.CgMethod
@@ -13,21 +12,18 @@ import org.utbot.framework.codegen.domain.models.CgMethodsCluster
 import org.utbot.framework.codegen.domain.models.CgRegion
 import org.utbot.framework.codegen.domain.models.CgStatement
 import org.utbot.framework.codegen.domain.models.CgStaticsRegion
-import org.utbot.framework.codegen.domain.models.CgVariable
-import org.utbot.framework.codegen.domain.models.SpringTestClassModel
-import org.utbot.framework.codegen.domain.models.builders.TypedModelWrappers
+import org.utbot.framework.codegen.domain.models.SimpleTestClassModel
+import org.utbot.framework.codegen.tree.fieldmanager.ClassFieldManagerFacade
 import org.utbot.framework.plugin.api.UtExecution
-import org.utbot.framework.plugin.api.UtSpringContextModel
 import org.utbot.framework.plugin.api.util.id
-import java.lang.Exception
 
 abstract class CgAbstractSpringTestClassConstructor(context: CgContext) :
-    CgAbstractTestClassConstructor<SpringTestClassModel>(context) {
+    CgAbstractTestClassConstructor<SimpleTestClassModel>(context) {
 
     protected val variableConstructor: CgSpringVariableConstructor =
         CgComponents.getVariableConstructorBy(context) as CgSpringVariableConstructor
 
-    override fun constructTestClassBody(testClassModel: SpringTestClassModel): CgClassBody {
+    override fun constructTestClassBody(testClassModel: SimpleTestClassModel): CgClassBody {
         return buildClassBody(currentTestClass) {
 
             // TODO: support inner classes here
@@ -80,7 +76,7 @@ abstract class CgAbstractSpringTestClassConstructor(context: CgContext) :
         return if (regions.any()) regions else null
     }
 
-    abstract fun constructClassFields(testClassModel: SpringTestClassModel): List<CgFieldDeclaration>
+    abstract fun constructClassFields(testClassModel: SimpleTestClassModel): List<CgFieldDeclaration>
 
     /**
      * Here "additional" means that these tests are not obtained from
@@ -90,49 +86,6 @@ abstract class CgAbstractSpringTestClassConstructor(context: CgContext) :
 
     open fun constructAdditionalUtilMethods(): CgMethodsCluster? = null
 
-    protected fun constructFieldsWithAnnotation(
-        fieldManager: CgClassFieldManager,
-        groupedModelsByClassId: TypedModelWrappers,
-    ): List<CgFieldDeclaration> {
-        val annotationClassId = fieldManager.annotationType
-        val annotation = addAnnotation(annotationClassId, Field)
-
-        val constructedDeclarations = mutableListOf<CgFieldDeclaration>()
-        for ((classId, modelWrappers) in groupedModelsByClassId) {
-
-            val modelWrapper = modelWrappers.firstOrNull() ?: continue
-            val model = modelWrapper.model
-
-            val fieldWithAnnotationIsRequired = fieldManager.fieldWithAnnotationIsRequired(model.classId)
-            if (!fieldWithAnnotationIsRequired) {
-                continue
-            }
-
-            val baseVarName = fieldManager.constructBaseVarName(model)
-
-            val createdVariable = variableConstructor.getOrCreateVariable(model, baseVarName) as? CgVariable
-                ?: error("`CgVariable` cannot be constructed from a $model model")
-
-            val declaration = CgDeclaration(classId, variableName = createdVariable.name, initializer = null)
-
-            constructedDeclarations += CgFieldDeclaration(
-                ownerClassId = currentTestClass,
-                declaration,
-                annotation
-            )
-
-            modelWrappers
-                .forEach { modelWrapper ->
-
-                    valueByUtModelWrapper[modelWrapper] = createdVariable
-
-                    variableConstructor.annotatedModelGroups
-                        .getOrPut(annotationClassId) { mutableSetOf() } += modelWrapper
-                }
-        }
-
-        return constructedDeclarations
-    }
 
     /**
      * Clears the results of variable instantiations that occurred
@@ -145,8 +98,7 @@ abstract class CgAbstractSpringTestClassConstructor(context: CgContext) :
      * but it will take very long time to do it now.
      */
     private fun clearUnwantedVariableModels() {
-        val trustedListOfModels =
-            variableConstructor.annotatedModelGroups.values.flatten() + listOf(UtSpringContextModel.wrap())
+        val trustedListOfModels = ClassFieldManagerFacade(context).findTrustedModels()
 
         valueByUtModelWrapper
             .filterNot { it.key in trustedListOfModels }