Introduce autowired class variables #2240 (#2248)

EgorkaKulikov · web-flow · commit 8227e23a0924 · 2023-06-16T13:41:37.000+03:00
diff --git a/utbot-framework/src/main/kotlin/org/utbot/framework/codegen/domain/models/TestClassModel.kt b/utbot-framework/src/main/kotlin/org/utbot/framework/codegen/domain/models/TestClassModel.kt
@@ -22,15 +22,17 @@ class SimpleTestClassModel(
 
 /**
  * Extended [SimpleTestClassModel] for Spring analysis reasons
- *
- * @param injectingMocksClass a class to inject other mocks into
- * @param mockedClasses variables of test class to represent mocked instances
  */
 class SpringTestClassModel(
     classUnderTest: ClassId,
     methodTestSets: List<CgMethodTestSet>,
     nestedClasses: List<SimpleTestClassModel>,
-    val thisInstanceModels: TypedModelWrappers = mapOf(),
-    val thisInstanceDependentMocks: TypedModelWrappers = mapOf(),
+    val springSpecificInformation: SpringSpecificInformation,
 ): TestClassModel(classUnderTest, methodTestSets, nestedClasses)
 
+
+class SpringSpecificInformation(
+    val thisInstanceModels: TypedModelWrappers = mapOf(),
+    val thisInstanceDependentMocks: TypedModelWrappers = mapOf(),
+    val applicationContextModels: TypedModelWrappers = mapOf(),
+)
diff --git a/utbot-framework/src/main/kotlin/org/utbot/framework/codegen/domain/models/builders/SpringTestClassModelBuilder.kt b/utbot-framework/src/main/kotlin/org/utbot/framework/codegen/domain/models/builders/SpringTestClassModelBuilder.kt
@@ -3,6 +3,7 @@ package org.utbot.framework.codegen.domain.models.builders
 import org.utbot.framework.codegen.domain.UtModelWrapper
 import org.utbot.framework.codegen.domain.context.CgContext
 import org.utbot.framework.codegen.domain.models.CgMethodTestSet
+import org.utbot.framework.codegen.domain.models.SpringSpecificInformation
 import org.utbot.framework.codegen.domain.models.SpringTestClassModel
 import org.utbot.framework.plugin.api.ClassId
 import org.utbot.framework.plugin.api.UtArrayModel
@@ -11,7 +12,6 @@ import org.utbot.framework.plugin.api.UtClassRefModel
 import org.utbot.framework.plugin.api.UtCompositeModel
 import org.utbot.framework.plugin.api.UtDirectSetFieldModel
 import org.utbot.framework.plugin.api.UtEnumConstantModel
-import org.utbot.framework.plugin.api.UtExecutableCallModel
 import org.utbot.framework.plugin.api.UtLambdaModel
 import org.utbot.framework.plugin.api.UtModel
 import org.utbot.framework.plugin.api.UtNullModel
@@ -27,20 +27,20 @@ class SpringTestClassModelBuilder(val context: CgContext): TestClassModelBuilder
 
     override fun createTestClassModel(classUnderTest: ClassId, testSets: List<CgMethodTestSet>): SpringTestClassModel {
         val baseModel = SimpleTestClassModelBuilder(context).createTestClassModel(classUnderTest, testSets)
-        val (thisInstanceModels, dependentMockModels) = collectThisInstanceAndDependentModels(testSets)
+        val springSpecificInformation = collectSpecificModelsForClassVariables(testSets)
 
         return SpringTestClassModel(
             classUnderTest = baseModel.classUnderTest,
             methodTestSets = baseModel.methodTestSets,
             nestedClasses = baseModel.nestedClasses,
-            thisInstanceModels = thisInstanceModels,
-            thisInstanceDependentMocks = dependentMockModels
+            springSpecificInformation = springSpecificInformation
         )
     }
 
-    private fun collectThisInstanceAndDependentModels(testSets: List<CgMethodTestSet>): Pair<TypedModelWrappers, TypedModelWrappers> {
+    private fun collectSpecificModelsForClassVariables(testSets: List<CgMethodTestSet>): SpringSpecificInformation {
         val thisInstanceModels = mutableSetOf<UtModelWrapper>()
         val thisInstancesDependentModels = mutableSetOf<UtModelWrapper>()
+        val stateBeforeDependentModels = mutableSetOf<UtModelWrapper>()
 
         with(context) {
             for ((testSetIndex, testSet) in testSets.withIndex()) {
@@ -51,8 +51,13 @@ class SpringTestClassModelBuilder(val context: CgContext): TestClassModelBuilder
                                 .filterNotNull()
                                 .forEach { model ->
                                     thisInstanceModels += model.wrap()
-                                    thisInstancesDependentModels += collectByThisInstanceModel(model)
+                                    thisInstancesDependentModels += collectByModel(model)
+
                                 }
+
+                            (execution.stateBefore.parameters + execution.stateBefore.thisInstance)
+                                .filterNotNull()
+                                .forEach { model -> stateBeforeDependentModels += collectByModel(model) }
                         }
                     }
                 }
@@ -65,10 +70,18 @@ class SpringTestClassModelBuilder(val context: CgContext): TestClassModelBuilder
                     cgModel.model.isMockModel() && cgModel !in thisInstanceModels
                 }
 
-        return thisInstanceModels.groupByClassId() to dependentMockModels.groupByClassId()
+        val applicationContextModels = stateBeforeDependentModels
+                .filter { it.model is UtSpringContextModel }
+                .toSet()
+
+        return SpringSpecificInformation(
+            thisInstanceModels.groupByClassId(),
+            dependentMockModels.groupByClassId(),
+            applicationContextModels.groupByClassId(),
+        )
     }
 
-    private fun collectByThisInstanceModel(model: UtModel): Set<UtModelWrapper> {
+    private fun collectByModel(model: UtModel): Set<UtModelWrapper> {
         val dependentModels = mutableSetOf<UtModelWrapper>()
         collectRecursively(model, dependentModels)
 
diff --git a/utbot-framework/src/main/kotlin/org/utbot/framework/codegen/tree/CgAbstractSpringTestClassConstructor.kt b/utbot-framework/src/main/kotlin/org/utbot/framework/codegen/tree/CgAbstractSpringTestClassConstructor.kt
@@ -85,17 +85,9 @@ abstract class CgAbstractSpringTestClassConstructor(context: CgContext):
     abstract fun constructAdditionalMethods(): CgMethodsCluster
 
     protected fun constructFieldsWithAnnotation(
+        annotationClassId: ClassId,
         groupedModelsByClassId: TypedModelWrappers,
-        annotationClassId: ClassId
     ): List<CgFieldDeclaration> {
-        require(
-            annotationClassId == injectMocksClassId ||
-                    annotationClassId == mockClassId ||
-                    annotationClassId == autowiredClassId
-        ) {
-            error("Unexpected annotation classId -- $annotationClassId")
-        }
-
         val annotation = statementConstructor.annotation(annotationClassId)
 
         val constructedDeclarations = mutableListOf<CgFieldDeclaration>()
@@ -111,10 +103,8 @@ abstract class CgAbstractSpringTestClassConstructor(context: CgContext):
                 valueByUtModelWrapper[key] = createdVariable
             }
 
-            when (annotationClassId) {
-                injectMocksClassId -> variableConstructor.injectedMocksModelsVariables += listOfUtModels
-                mockClassId -> variableConstructor.mockedModelsVariables += listOfUtModels
-            }
+            variableConstructor.annotatedModelVariables
+                .getOrPut(annotationClassId) { mutableSetOf() } += listOfUtModels
         }
 
         return constructedDeclarations
@@ -130,12 +120,8 @@ abstract class CgAbstractSpringTestClassConstructor(context: CgContext):
      * related side effects and just creating a variable definition,
      * but it will take very long time to do it now.
      */
-    protected fun clearUnwantedVariableModels() {
-        val whiteListOfModels =
-            listOf(
-                variableConstructor.mockedModelsVariables,
-                variableConstructor.injectedMocksModelsVariables
-            ).flatten()
+    private fun clearUnwantedVariableModels() {
+        val whiteListOfModels = variableConstructor.annotatedModelVariables.values.flatten()
 
         valueByUtModelWrapper
             .filter { it.key !in whiteListOfModels }
diff --git a/utbot-framework/src/main/kotlin/org/utbot/framework/codegen/tree/CgSpringIntegrationTestClassConstructor.kt b/utbot-framework/src/main/kotlin/org/utbot/framework/codegen/tree/CgSpringIntegrationTestClassConstructor.kt
@@ -9,7 +9,8 @@ import org.utbot.framework.codegen.domain.models.SpringTestClassModel
 class CgSpringIntegrationTestClassConstructor(context: CgContext) : CgAbstractSpringTestClassConstructor(context) {
 
     override fun constructClassFields(testClassModel: SpringTestClassModel): List<CgFieldDeclaration> {
-        return constructFieldsWithAnnotation(testClassModel.thisInstanceModels, autowiredClassId)
+        val applicationContextModels = testClassModel.springSpecificInformation.applicationContextModels
+        return constructFieldsWithAnnotation(autowiredClassId, applicationContextModels)
     }
 
     override fun constructAdditionalMethods() = CgMethodsCluster(
diff --git a/utbot-framework/src/main/kotlin/org/utbot/framework/codegen/tree/CgSpringUnitTestClassConstructor.kt b/utbot-framework/src/main/kotlin/org/utbot/framework/codegen/tree/CgSpringUnitTestClassConstructor.kt
@@ -27,10 +27,12 @@ class CgSpringUnitTestClassConstructor(context: CgContext) : CgAbstractSpringTes
 
     override fun constructClassFields(testClassModel: SpringTestClassModel): List<CgFieldDeclaration> {
         val fields = mutableListOf<CgFieldDeclaration>()
+        val thisInstances = testClassModel.springSpecificInformation.thisInstanceModels
+        val mocks = testClassModel.springSpecificInformation.thisInstanceDependentMocks
 
-        if (testClassModel.thisInstanceDependentMocks.isNotEmpty()) {
-            val mockedFields = constructFieldsWithAnnotation(testClassModel.thisInstanceDependentMocks, mockClassId)
-            val injectingMocksFields = constructFieldsWithAnnotation(testClassModel.thisInstanceModels, injectMocksClassId)
+        if (mocks.isNotEmpty()) {
+            val mockedFields = constructFieldsWithAnnotation(mockClassId, mocks)
+            val injectingMocksFields = constructFieldsWithAnnotation(injectMocksClassId, thisInstances)
 
             fields += injectingMocksFields
             fields += mockedFields
diff --git a/utbot-framework/src/main/kotlin/org/utbot/framework/codegen/tree/CgSpringVariableConstructor.kt b/utbot-framework/src/main/kotlin/org/utbot/framework/codegen/tree/CgSpringVariableConstructor.kt
@@ -1,26 +1,24 @@
 package org.utbot.framework.codegen.tree
 
 import org.utbot.framework.codegen.domain.UtModelWrapper
+import org.utbot.framework.codegen.domain.builtin.autowiredClassId
+import org.utbot.framework.codegen.domain.builtin.injectMocksClassId
+import org.utbot.framework.codegen.domain.builtin.mockClassId
 import org.utbot.framework.codegen.domain.context.CgContext
 import org.utbot.framework.codegen.domain.models.CgValue
 import org.utbot.framework.codegen.domain.models.CgVariable
+import org.utbot.framework.plugin.api.ClassId
 import org.utbot.framework.plugin.api.UtAssembleModel
 import org.utbot.framework.plugin.api.UtCompositeModel
 import org.utbot.framework.plugin.api.UtModel
 import org.utbot.framework.plugin.api.UtSpringContextModel
 import org.utbot.framework.plugin.api.isMockModel
 
 class CgSpringVariableConstructor(context: CgContext) : CgVariableConstructor(context) {
-    val injectedMocksModelsVariables: MutableSet<UtModelWrapper> = mutableSetOf()
-    val mockedModelsVariables: MutableSet<UtModelWrapper> = mutableSetOf()
+    val annotatedModelVariables: MutableMap<ClassId, MutableSet<UtModelWrapper>> = mutableMapOf()
 
     override fun getOrCreateVariable(model: UtModel, name: String?): CgValue {
-        if (model is UtSpringContextModel) {
-            // TODO also, properly render corresponding @Autowired field(s), and make sure name isn't taken
-            return CgVariable(model.modelName, model.classId)
-        }
-
-        val alreadyCreatedInjectMocks = findCgValueByModel(model, injectedMocksModelsVariables)
+        val alreadyCreatedInjectMocks = findCgValueByModel(model, annotatedModelVariables[injectMocksClassId])
         if (alreadyCreatedInjectMocks != null) {
             val fields: Collection<UtModel> = when (model) {
                 is UtCompositeModel -> model.fields.values
@@ -33,7 +31,7 @@ class CgSpringVariableConstructor(context: CgContext) : CgVariableConstructor(co
             return alreadyCreatedInjectMocks
         }
 
-        val alreadyCreatedMock = findCgValueByModel(model, mockedModelsVariables)
+        val alreadyCreatedMock = findCgValueByModel(model, annotatedModelVariables[mockClassId])
         if (alreadyCreatedMock != null) {
             if (model.isMockModel()) {
                 mockFrameworkManager.createMockForVariable(
@@ -45,11 +43,26 @@ class CgSpringVariableConstructor(context: CgContext) : CgVariableConstructor(co
             return alreadyCreatedMock
         }
 
-        return super.getOrCreateVariable(model, name)
+        val alreadyCreatedAutowired = findCgValueByModel(model, annotatedModelVariables[autowiredClassId])
+        if (alreadyCreatedAutowired != null) {
+            return alreadyCreatedAutowired
+        }
+
+        return when (model) {
+            is UtSpringContextModel -> constructSpringContext(model, name)
+            else -> super.getOrCreateVariable(model, name)
+        }
+    }
+
+    private fun constructSpringContext(model: UtSpringContextModel, baseName: String?): CgValue {
+        val obj = newVar(model.classId, baseName) { utilsClassId[createInstance](model.classId.name) }
+
+        valueByUtModelWrapper[model.wrap()] = obj
+        return obj
     }
 
-    private fun findCgValueByModel(model: UtModel, setOfModels: Set<UtModelWrapper>): CgValue? {
-        val key = setOfModels.find { it == model.wrap() }
+    private fun findCgValueByModel(model: UtModel, setOfModels: Set<UtModelWrapper>?): CgValue? {
+        val key = setOfModels?.find { it == model.wrap() } ?: return null
         return valueByUtModelWrapper[key]
     }
 }
diff --git a/utbot-framework/src/main/kotlin/org/utbot/framework/codegen/tree/CgVariableConstructor.kt b/utbot-framework/src/main/kotlin/org/utbot/framework/codegen/tree/CgVariableConstructor.kt
@@ -52,6 +52,7 @@ 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.UtReferenceModel
+import org.utbot.framework.plugin.api.UtSpringContextModel
 import org.utbot.framework.plugin.api.UtStatementCallModel
 import org.utbot.framework.plugin.api.UtVoidModel
 import org.utbot.framework.plugin.api.util.classClassId

Original file line number	Diff line number	Diff line change
`@@ -9,7 +9,8 @@ import org.utbot.framework.codegen.domain.models.SpringTestClassModel`
`9`	`9`	`class CgSpringIntegrationTestClassConstructor(context: CgContext) : CgAbstractSpringTestClassConstructor(context) {`
`10`	`10`
`11`	`11`	`override fun constructClassFields(testClassModel: SpringTestClassModel): List<CgFieldDeclaration> {`
`12`		`- return constructFieldsWithAnnotation(testClassModel.thisInstanceModels, autowiredClassId)`
	`12`	`+ val applicationContextModels = testClassModel.springSpecificInformation.applicationContextModels`
	`13`	`+ return constructFieldsWithAnnotation(autowiredClassId, applicationContextModels)`
`13`	`14`	`}`
`14`	`15`
`15`	`16`	`override fun constructAdditionalMethods() = CgMethodsCluster(`