Support autowiring collections in Spring Unit tests

utbot-framework-api/src/main/kotlin/org/utbot/framework/plugin/api/Api.kt

@@ -398,6 +398,18 @@ fun UtModel.hasDefaultValue() =
  */
 fun UtModel.isMockModel() = this is UtCompositeModel && isMock
 
+/**
+ * Checks that this [UtModel] must be constructed with @Spy annotation in generated tests.
+ * Used only for construct variables with @Spy annotation.
+ */
+fun UtModel.canBeSpied(): Boolean {
+    val javaClass = this.classId.jClass
+
+    return this is UtAssembleModel &&
+            (Collection::class.java.isAssignableFrom(javaClass)
+                    || Map::class.java.isAssignableFrom(javaClass))
+}
+
 /**
  * Get model id (symbolic null value for UtNullModel)
  * or null if model has no id (e.g., a primitive model) or the id is null.

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

@@ -757,12 +757,9 @@ object SpringBoot : DependencyInjectionFramework(
  *
  * Used as a key in [valueByUtModelWrapper].
  * Was introduced primarily for shared among all test methods global variables.
- *
- * `modelTagName` is used to distinguish between variables with annotation @Mock that have the same model
  */
 data class UtModelWrapper(
     val testSetId: Int,
     val executionId: Int,
     val model: UtModel,
-    val modelTagName: String?
 )

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

@@ -25,6 +25,11 @@ internal val mockClassId = BuiltinClassId(
     simpleName = "Mock",
 )
 
+internal val spyClassId = BuiltinClassId(
+    canonicalName = "org.mockito.Spy",
+    simpleName = "Spy"
+)
+
 internal val injectMocksClassId = BuiltinClassId(
     canonicalName = "org.mockito.InjectMocks",
     simpleName = "InjectMocks",

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

@@ -318,6 +318,26 @@ val closeMethodId = MethodId(
     parameters = emptyList(),
 )
 
+private val clearCollectionMethodId = MethodId(
+    classId = Collection::class.java.id,
+    name = "clear",
+    returnType = voidClassId,
+    parameters = emptyList()
+)
+
+private val clearMapMethodId = MethodId(
+    classId = Map::class.java.id,
+    name = "clear",
+    returnType = voidClassId,
+    parameters = emptyList()
+)
+
+fun clearMethodId(javaClass: Class<*>): MethodId = when {
+    Collection::class.java.isAssignableFrom(javaClass) -> clearCollectionMethodId
+    Map::class.java.isAssignableFrom(javaClass) -> clearMapMethodId
+    else -> error("Clear method is not implemented for $javaClass")
+}
+
 val mocksAutoCloseable: Set<ClassId> = setOf(
     MockitoStaticMocking.mockedStaticClassId,
     MockitoStaticMocking.mockedConstructionClassId

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

@@ -457,12 +457,11 @@ interface CgContextOwner {
     val getLambdaMethod: MethodId
         get() = utilMethodProvider.getLambdaMethodMethodId
 
-    fun UtModel.wrap(modelTagName: String? = null): UtModelWrapper =
+    fun UtModel.wrap(): UtModelWrapper =
         UtModelWrapper(
             testSetId = currentTestSetId,
             executionId = currentExecutionId,
             model = this,
-            modelTagName = modelTagName
         )
 }
 

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

@@ -33,5 +33,6 @@ class SpringTestClassModel(
 class SpringSpecificInformation(
     val thisInstanceModels: TypedModelWrappers,
     val thisInstanceDependentMocks: TypedModelWrappers,
+    val thisInstanceDependentSpies: TypedModelWrappers,
     val autowiredFromContextModels: TypedModelWrappers,
 )

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

@@ -12,20 +12,19 @@ import org.utbot.framework.plugin.api.UtAssembleModel
 import org.utbot.framework.plugin.api.UtClassRefModel
 import org.utbot.framework.plugin.api.UtCompositeModel
 import org.utbot.framework.plugin.api.UtCustomModel
-import org.utbot.framework.plugin.api.UtDirectSetFieldModel
 import org.utbot.framework.plugin.api.UtEnumConstantModel
 import org.utbot.framework.plugin.api.UtLambdaModel
 import org.utbot.framework.plugin.api.UtModel
 import org.utbot.framework.plugin.api.UtNullModel
 import org.utbot.framework.plugin.api.UtPrimitiveModel
-import org.utbot.framework.plugin.api.UtStatementCallModel
 import org.utbot.framework.plugin.api.UtVoidModel
 import org.utbot.framework.plugin.api.isMockModel
 import org.utbot.framework.plugin.api.util.SpringModelUtils.isAutowiredFromContext
+import org.utbot.framework.plugin.api.canBeSpied
 
 typealias TypedModelWrappers = Map<ClassId, Set<UtModelWrapper>>
 
-class SpringTestClassModelBuilder(val context: CgContext):
+class SpringTestClassModelBuilder(val context: CgContext) :
     TestClassModelBuilder(),
     CgContextOwner by context {
 
@@ -54,13 +53,13 @@ class SpringTestClassModelBuilder(val context: CgContext):
                             .filterNotNull()
                             .forEach { model ->
                                 thisInstanceModels += model.wrap()
-                                thisInstancesDependentModels += collectByModel(model)
+                                thisInstancesDependentModels += collectDependentModels(model)
 
                             }
 
                         (execution.stateBefore.parameters + execution.stateBefore.thisInstance)
                             .filterNotNull()
-                            .forEach { model -> stateBeforeDependentModels += collectByModel(model) }
+                            .forEach { model -> stateBeforeDependentModels += collectDependentModels(model) }
                     }
                 }
             }
@@ -72,78 +71,66 @@ class SpringTestClassModelBuilder(val context: CgContext):
                     cgModel.model.isMockModel() && cgModel !in thisInstanceModels
                 }
 
+        val dependentSpyModels =
+            thisInstancesDependentModels
+                .filterTo(mutableSetOf()) { cgModel ->
+                    cgModel.model.canBeSpied() && cgModel !in thisInstanceModels && cgModel !in dependentMockModels
+                }
+
         val autowiredFromContextModels =
             stateBeforeDependentModels.filterTo(HashSet()) { it.model.isAutowiredFromContext() }
 
         return SpringSpecificInformation(
             thisInstanceModels.groupByClassId(),
             dependentMockModels.groupByClassId(),
+            dependentSpyModels.groupByClassId(),
             autowiredFromContextModels.groupByClassId(),
         )
     }
 
-    private fun collectByModel(model: UtModel): Set<UtModelWrapper> {
+    private fun collectDependentModels(model: UtModel): Set<UtModelWrapper> {
         val dependentModels = mutableSetOf<UtModelWrapper>()
 
-        collectRecursively(model.wrap(), dependentModels)
-
-        return dependentModels
-    }
-
-    private fun Set<UtModelWrapper>.groupByClassId(): TypedModelWrappers {
-        val classModels = mutableMapOf<ClassId, Set<UtModelWrapper>>()
-
-        for (modelGroup in this.groupBy { it.model.classId }) {
-            classModels[modelGroup.key] = modelGroup.value.toSet()
-        }
-
-        return classModels
-    }
-
-    private fun collectRecursively(currentModelWrapper: UtModelWrapper, allModels: MutableSet<UtModelWrapper>) {
-        if (!allModels.add(currentModelWrapper)) {
-            return
-        }
-
-        when (val currentModel = currentModelWrapper.model) {
+        when (model) {
             is UtNullModel,
             is UtPrimitiveModel,
             is UtClassRefModel,
             is UtVoidModel,
             is UtEnumConstantModel,
-            is UtCustomModel-> {}
+            is UtCustomModel -> {}
             is UtLambdaModel -> {
-                currentModel.capturedValues.forEach { collectRecursively(it.wrap(), allModels) }
+                model.capturedValues.forEach { dependentModels.add(it.wrap()) }
             }
             is UtArrayModel -> {
-                currentModel.stores.values.forEach { collectRecursively(it.wrap(), allModels) }
-                if (currentModel.stores.count() < currentModel.length) {
-                    collectRecursively(currentModel.constModel.wrap(), allModels)
+                model.stores.values.forEach { dependentModels.add(it.wrap()) }
+                if (model.stores.count() < model.length) {
+                    dependentModels.add(model.constModel.wrap())
                 }
             }
             is UtCompositeModel -> {
                 // Here we traverse fields only.
                 // Traversing mocks as well will result in wrong models playing
                 // a role of class fields with @Mock annotation.
-                currentModel.fields.forEach { (fieldId, model) ->
-                    // We use `modelTagName` in order to distinguish mock models
-                    val modeTagName = if(model.isMockModel()) fieldId.name else null
-                    collectRecursively(model.wrap(modeTagName), allModels)
-                }
+                model.fields.forEach { (_, model) -> dependentModels.add(model.wrap()) }
             }
             is UtAssembleModel -> {
-                currentModel.instantiationCall.instance?.let { collectRecursively(it.wrap(), allModels) }
-                currentModel.instantiationCall.params.forEach { collectRecursively(it.wrap(), allModels) }
-
-                currentModel.modificationsChain.forEach { stmt ->
-                    stmt.instance?.let { collectRecursively(it.wrap(), allModels) }
-                    when (stmt) {
-                        is UtStatementCallModel -> stmt.params.forEach { collectRecursively(it.wrap(), allModels) }
-                        is UtDirectSetFieldModel -> collectRecursively(stmt.fieldModel.wrap(), allModels)
-                    }
-                }
+                model.instantiationCall.instance?.let { dependentModels.add(it.wrap()) }
+                model.instantiationCall.params.forEach { dependentModels.add(it.wrap()) }
             }
-            //Python, JavaScript, Go models are not required in Spring
         }
+
+        return dependentModels
+    }
+
+
+    private fun Set<UtModelWrapper>.groupByClassId(): TypedModelWrappers {
+        val classModels = mutableMapOf<ClassId, Set<UtModelWrapper>>()
+
+        for (modelGroup in this.groupBy { it.model.classId }) {
+            classModels[modelGroup.key] = modelGroup.value.toSet()
+        }
+
+        return classModels
     }
+
 }

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/services/framework/SpyFrameworkManager.kt

@@ -0,0 +1,12 @@
+package org.utbot.framework.codegen.services.framework
+
+import org.utbot.framework.codegen.domain.context.CgContext
+import org.utbot.framework.plugin.api.UtAssembleModel
+
+class SpyFrameworkManager(context: CgContext) : CgVariableConstructorComponent(context) {
+
+    fun spyForVariable(model: UtAssembleModel){
+        variableConstructor.constructAssembleForVariable(model)
+    }
+
+}

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

@@ -1,6 +1,5 @@
 package org.utbot.framework.codegen.tree
 
-import org.utbot.framework.codegen.domain.UtModelWrapper
 import org.utbot.framework.codegen.domain.builtin.TestClassUtilMethodProvider
 import org.utbot.framework.codegen.domain.context.CgContext
 import org.utbot.framework.codegen.domain.models.AnnotationTarget.*
@@ -19,7 +18,6 @@ import org.utbot.framework.codegen.domain.models.SpringTestClassModel
 import org.utbot.framework.codegen.domain.models.builders.TypedModelWrappers
 import org.utbot.framework.plugin.api.UtExecution
 import org.utbot.framework.plugin.api.UtSpringContextModel
-import org.utbot.framework.plugin.api.util.SpringModelUtils.getBeanNameOrNull
 import org.utbot.framework.plugin.api.util.id
 import java.lang.Exception
 
@@ -102,18 +100,18 @@ abstract class CgAbstractSpringTestClassConstructor(context: CgContext) :
         val constructedDeclarations = mutableListOf<CgFieldDeclaration>()
         for ((classId, modelWrappers) in groupedModelsByClassId) {
 
-            val fieldWithAnnotationIsRequired = fieldManager.fieldWithAnnotationIsRequired(modelWrappers)
+            val modelWrapper = modelWrappers.firstOrNull() ?: continue
+            val model = modelWrapper.model
+
+            val fieldWithAnnotationIsRequired = fieldManager.fieldWithAnnotationIsRequired(model.classId)
             if (!fieldWithAnnotationIsRequired) {
                 continue
             }
 
-            val modelWrapper = modelWrappers.firstOrNull() ?: continue
-            val model = modelWrapper.model
-
-            val baseVarName = model.getBeanNameOrNull()
+            val baseVarName = fieldManager.constructBaseVarName(model)
 
             val createdVariable = variableConstructor.getOrCreateVariable(model, baseVarName) as? CgVariable
-                ?: error("`UtCompositeModel` model was expected, but $model was found")
+                ?: error("`CgVariable` cannot be constructed from a $model model")
 
             val declaration = CgDeclaration(classId, variableName = createdVariable.name, initializer = null)
 
@@ -126,17 +124,10 @@ abstract class CgAbstractSpringTestClassConstructor(context: CgContext) :
             modelWrappers
                 .forEach { modelWrapper ->
 
-                    val modelWrapperWithNullTagName = UtModelWrapper(
-                        testSetId = modelWrapper.testSetId,
-                        executionId = modelWrapper.executionId,
-                        model = modelWrapper.model,
-                        modelTagName = null,
-                    )
-
-                    valueByUtModelWrapper[modelWrapperWithNullTagName] = createdVariable
+                    valueByUtModelWrapper[modelWrapper] = createdVariable
 
                     variableConstructor.annotatedModelGroups
-                        .getOrPut(annotationClassId) { mutableSetOf() } += modelWrapperWithNullTagName
+                        .getOrPut(annotationClassId) { mutableSetOf() } += modelWrapper
                 }
         }