Corrects collecting annotated field models (#2513)

Author: EgorkaKulikov

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

@@ -27,9 +27,10 @@ class CgSpringUnitTestClassConstructor(context: CgContext) : CgAbstractSpringTes
     private lateinit var mockitoCloseableVariable: CgValue
     private lateinit var spyClearVariables: List<CgValue>
 
-    private val injectingMocksFieldsManager = CgInjectingMocksFieldsManager(context)
     private val mocksFieldsManager = CgMockedFieldsManager(context)
     private val spiesFieldsManager = CgSpiedFieldsManager(context)
+    private val injectingMocksFieldsManager =
+        CgInjectingMocksFieldsManager(context, mocksFieldsManager, spiesFieldsManager)
 
     override fun constructClassFields(testClassModel: SimpleTestClassModel): List<CgFieldDeclaration> {
         val fields = mutableListOf<CgFieldDeclaration>()

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

@@ -18,11 +18,9 @@ abstract class CgAbstractClassFieldManager(context: CgContext) :
     CgClassFieldManager,
     CgContextOwner by context {
 
-    protected val annotatedModelGroups: MutableMap<ClassId, MutableSet<UtModelWrapper>> = mutableMapOf()
+    val annotatedModels: MutableSet<UtModelWrapper> = mutableSetOf()
     protected val modelGroupsProvider = ModelGroupsProvider(context)
 
-    fun getManagedModels(): Set<UtModelWrapper> = annotatedModelGroups[annotationType] ?: emptySet()
-
     fun findCgValueByModel(model: UtModel, setOfModels: Set<UtModelWrapper>?): CgValue? {
         val key = setOfModels?.find { it == model.wrap() } ?: return null
         return valueByUtModelWrapper[key]
@@ -60,7 +58,7 @@ abstract class CgAbstractClassFieldManager(context: CgContext) :
 
             modelWrappers.forEach { modelWrapper ->
                 valueByUtModelWrapper[modelWrapper] = createdVariable
-                annotatedModelGroups.getOrPut(annotationType) { mutableSetOf() } += modelWrapper
+                annotatedModels += modelWrapper
             }
         }
 

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

@@ -13,7 +13,11 @@ import org.utbot.framework.plugin.api.UtModel
 import org.utbot.framework.plugin.api.UtModelWithCompositeOrigin
 import org.utbot.framework.plugin.api.isMockModel
 
-class CgInjectingMocksFieldsManager(val context: CgContext) : CgAbstractClassFieldManager(context) {
+class CgInjectingMocksFieldsManager(
+    val context: CgContext,
+    private val mocksFieldsManager: CgMockedFieldsManager,
+    private val spiesFieldsManager: CgSpiedFieldsManager,
+    ) : CgAbstractClassFieldManager(context) {
     init {
         relevantFieldManagers += this
     }
@@ -37,10 +41,10 @@ class CgInjectingMocksFieldsManager(val context: CgContext) : CgAbstractClassFie
             val variableForField = variableConstructor.getOrCreateVariable(fieldModel)
 
             // is variable mocked by @Mock annotation
-            val isMocked = findCgValueByModel(fieldModel, annotatedModelGroups[mockClassId]) != null
+            val isMocked = findCgValueByModel(fieldModel, mocksFieldsManager.annotatedModels) != null
 
             // is variable spied by @Spy annotation
-            val isSpied = findCgValueByModel(fieldModel, annotatedModelGroups[spyClassId]) != null
+            val isSpied = findCgValueByModel(fieldModel, spiesFieldsManager.annotatedModels) != null
 
             // If field model is a mock model and is mocked by @Mock annotation in classFields or is spied by @Spy annotation,
             // it is set in the connected with instance under test automatically via @InjectMocks.

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

@@ -11,9 +11,7 @@ class ClassFieldManagerFacade(context: CgContext) : CgContextOwner by context {
 
     fun constructVariableForField(model: UtModel): CgValue? {
         relevantFieldManagers.forEach { manager ->
-            val managedModels = manager.getManagedModels()
-
-            val alreadyCreatedVariable = manager.findCgValueByModel(model, managedModels)
+            val alreadyCreatedVariable = manager.findCgValueByModel(model, manager.annotatedModels)
             if (alreadyCreatedVariable != null) {
                 manager.useVariableForModel(model, alreadyCreatedVariable)
                 return alreadyCreatedVariable
@@ -27,7 +25,7 @@ class ClassFieldManagerFacade(context: CgContext) : CgContextOwner by context {
         val trustedModels = mutableListOf(UtSpringContextModel.wrap())
 
          relevantFieldManagers.forEach { manager ->
-            trustedModels += manager.getManagedModels()
+            trustedModels += manager.annotatedModels
         }
 
         return trustedModels