Add support for class fields rendering with annotations

Author: sofurihafe

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

@@ -20,6 +20,16 @@ internal val mockitoClassId = BuiltinClassId(
     simpleName = "Mockito",
 )
 
+internal val mockClassId = BuiltinClassId(
+    canonicalName = "org.mockito.Mock",
+    simpleName = "Mock",
+)
+
+internal val injectMocksClassId = BuiltinClassId(
+    canonicalName = "org.mockito.InjectMocks",
+    simpleName = "InjectMocks",
+)
+
 internal val ongoingStubbingClassId = BuiltinClassId(
     canonicalName = "org.mockito.stubbing.OngoingStubbing",
     simpleName = "OngoingStubbing",

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

@@ -82,6 +82,7 @@ interface CgElement {
             is CgBreakStatement -> visit(element)
             is CgContinueStatement -> visit(element)
             is CgDeclaration -> visit(element)
+            is CgFieldDeclaration -> visit(element)
             is CgAssignment -> visit(element)
             is CgTypeCast -> visit(element)
             is CgIsInstance -> visit(element)
@@ -156,8 +157,19 @@ class CgClassBody(
     val methodRegions: List<CgMethodsCluster>,
     val staticDeclarationRegions: List<CgStaticsRegion>,
     val nestedClassRegions: List<CgNestedClassesRegion<*>>,
-    //TODO: use [CgFieldDeclaration] after PR-1788 merge
-    val fields: List<CgDeclaration> = emptyList(),
+    val fields: List<CgFieldDeclaration> = emptyList(),
+) : CgElement
+
+/**
+ * Field of a class.
+ * @property classId [ClassId] of the field owner class.
+ * @property declaration declaration itself.
+ * @property annotation optional annotation.
+ */
+class CgFieldDeclaration(
+    val declaration: CgDeclaration,
+    val annotation: CgAnnotation? = null,
+    val classId: ClassId? = null,
 ) : CgElement
 
 /**

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

@@ -1,6 +1,7 @@
 package org.utbot.framework.codegen.domain.models
 
 import org.utbot.framework.plugin.api.ClassId
+import org.utbot.framework.plugin.api.UtModel
 
 /**
  * Stores method test sets in a structure that replicates structure of their methods in [classUnderTest].
@@ -13,7 +14,7 @@ abstract class TestClassModel(
     val nestedClasses: List<SimpleTestClassModel>,
 )
 
-open class SimpleTestClassModel(
+class SimpleTestClassModel(
     classUnderTest: ClassId,
     methodTestSets: List<CgMethodTestSet>,
     nestedClasses: List<SimpleTestClassModel> = listOf(),
@@ -29,7 +30,7 @@ class SpringTestClassModel(
     classUnderTest: ClassId,
     methodTestSets: List<CgMethodTestSet>,
     nestedClasses: List<SimpleTestClassModel>,
-    val injectingMocksClass: ClassId? = null,
-    val mockedClasses: Set<ClassId> = setOf(),
+    val injectedMockModels: Map<ClassId, List<UtModel>> = mapOf(),
+    val mockedModels: Map<ClassId, List<UtModel>> = mapOf(),
 ): TestClassModel(classUnderTest, methodTestSets, nestedClasses)
 

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

@@ -21,27 +21,38 @@ class SpringTestClassModelBuilder: TestClassModelBuilder() {
 
     override fun createTestClassModel(classUnderTest: ClassId, testSets: List<CgMethodTestSet>): SpringTestClassModel {
         val baseModel = SimpleTestClassModelBuilder().createTestClassModel(classUnderTest, testSets)
-        val mockedClasses = collectMockedClassIds(classUnderTest, testSets)
+        val (injectedModels, mockedModels) = collectInjectedAndMockedModels(testSets)
 
         return SpringTestClassModel(
             baseModel.classUnderTest,
             baseModel.methodTestSets,
             baseModel.nestedClasses,
-            classUnderTest,
-            mockedClasses,
+            groupUtModelsByClassId(injectedModels),
+            groupUtModelsByClassId(mockedModels),
         )
     }
 
-    private fun collectMockedClassIds(
-        classUnderTest: ClassId,
+    private fun groupUtModelsByClassId(models: List<UtModel>): Map<ClassId, List<UtModel>> {
+        return models.groupBy { it.classId }
+    }
+
+    private fun collectInjectedAndMockedModels(
         testSets: List<CgMethodTestSet>,
-    ): Set<ClassId> {
+    ): Pair<List<UtModel>, List<UtModel>> {
+        val thisInstances = mutableListOf<UtModel>()
         val allModelsInExecution = mutableListOf<UtModel>()
 
         for (testSet in testSets) {
             for (execution in testSet.executions) {
-                execution.stateBefore.thisInstance?.let { allModelsInExecution += it }
-                execution.stateAfter.thisInstance?.let { allModelsInExecution += it }
+                execution.stateBefore.thisInstance?.let {
+                    allModelsInExecution += it
+                    thisInstances += it
+                }
+
+                execution.stateAfter.thisInstance?.let {
+                    allModelsInExecution += it
+                    thisInstances += it
+                }
 
                 allModelsInExecution += execution.stateBefore.parameters
                 allModelsInExecution += execution.stateAfter.parameters
@@ -53,16 +64,17 @@ class SpringTestClassModelBuilder: TestClassModelBuilder() {
         val allConstructedModels = mutableSetOf<UtModel>()
         allModelsInExecution.forEach { model -> collectRecursively(model, allConstructedModels) }
 
-        return allConstructedModels
+        val injectedModels = thisInstances.filterIsInstance<UtCompositeModel>()
+        val mockedModels = allConstructedModels
             .filter { it.isMockComposite() || it.isMockAssemble() }
-            .map { it.classId }
-            .filter {  it != classUnderTest }
-            .toSet()
+            .filter { it !in injectedModels }
+            .distinct()
 
+        return injectedModels to mockedModels
     }
 
     private fun collectRecursively(currentModel: UtModel, allModels: MutableSet<UtModel>) {
-        if (currentModel in allModels) {
+        if (allModels.any { it === currentModel }) {
             return
         }
 

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/renderer/CgAbstractRenderer.kt

@@ -41,6 +41,7 @@ import org.utbot.framework.codegen.domain.models.CgExecutableCall
 import org.utbot.framework.codegen.domain.models.CgMethodsCluster
 import org.utbot.framework.codegen.domain.models.CgExpression
 import org.utbot.framework.codegen.domain.models.CgFieldAccess
+import org.utbot.framework.codegen.domain.models.CgFieldDeclaration
 import org.utbot.framework.codegen.domain.models.CgForEachLoop
 import org.utbot.framework.codegen.domain.models.CgForLoop
 import org.utbot.framework.codegen.domain.models.CgFrameworkUtilMethod
@@ -571,6 +572,14 @@ abstract class CgAbstractRenderer(
         println(statementEnding)
     }
 
+    // Class field declaration
+
+    override fun visit(element: CgFieldDeclaration) {
+        element.annotation?.accept(this)
+        renderVisibility(element.declaration.variableType)
+        element.declaration.accept(this)
+    }
+
     // Variable assignment
 
     override fun visit(element: CgAssignment) {
@@ -888,7 +897,7 @@ abstract class CgAbstractRenderer(
         }
     }
 
-    protected abstract fun renderClassVisibility(classId: ClassId)
+    protected abstract fun renderVisibility(classId: ClassId)
 
     protected abstract fun renderClassModality(aClass: CgClass)
 

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/renderer/CgJavaRenderer.kt

@@ -32,6 +32,7 @@ import org.utbot.framework.codegen.domain.models.CgSwitchCase
 import org.utbot.framework.codegen.domain.models.CgSwitchCaseLabel
 import org.utbot.framework.codegen.domain.models.CgClass
 import org.utbot.framework.codegen.domain.models.CgClassBody
+import org.utbot.framework.codegen.domain.models.CgFieldDeclaration
 import org.utbot.framework.codegen.domain.models.CgFormattedString
 import org.utbot.framework.codegen.domain.models.CgFrameworkUtilMethod
 import org.utbot.framework.codegen.domain.models.CgLiteral
@@ -70,7 +71,7 @@ internal class CgJavaRenderer(context: CgRendererContext, printer: CgPrinter = C
             annotation.accept(this)
         }
 
-        renderClassVisibility(element.id)
+        renderVisibility(element.id)
         renderClassModality(element)
         if (element.isStatic) {
             print("static ")
@@ -92,6 +93,12 @@ internal class CgJavaRenderer(context: CgRendererContext, printer: CgPrinter = C
     }
 
     override fun visit(element: CgClassBody) {
+        // render class fields
+        for (field in element.fields) {
+            field.accept(this)
+            println()
+        }
+
         // render regions for test methods and utils
         val allRegions = element.methodRegions + element.nestedClassRegions + element.staticDeclarationRegions
         for ((i, region) in allRegions.withIndex()) {
@@ -101,6 +108,12 @@ internal class CgJavaRenderer(context: CgRendererContext, printer: CgPrinter = C
         }
     }
 
+    override fun visit(element: CgFieldDeclaration) {
+        element.annotation?.accept(this)
+        renderVisibility(element.declaration.variableType)
+        element.declaration.accept(this)
+    }
+
     override fun visit(element: CgArrayAnnotationArgument) {
         print("{")
         element.values.renderSeparated()
@@ -386,7 +399,7 @@ internal class CgJavaRenderer(context: CgRendererContext, printer: CgPrinter = C
 
     override fun escapeNamePossibleKeywordImpl(s: String): String = s
 
-    override fun renderClassVisibility(classId: ClassId) {
+    override fun renderVisibility(classId: ClassId) {
         when {
             classId.isPublic -> print("public ")
             classId.isProtected -> print("protected ")

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/renderer/CgKotlinRenderer.kt

@@ -38,6 +38,7 @@ import org.utbot.framework.codegen.domain.models.CgSwitchCase
 import org.utbot.framework.codegen.domain.models.CgSwitchCaseLabel
 import org.utbot.framework.codegen.domain.models.CgClass
 import org.utbot.framework.codegen.domain.models.CgClassBody
+import org.utbot.framework.codegen.domain.models.CgFieldDeclaration
 import org.utbot.framework.codegen.domain.models.CgFormattedString
 import org.utbot.framework.codegen.domain.models.CgFrameworkUtilMethod
 import org.utbot.framework.codegen.domain.models.CgLiteral
@@ -85,7 +86,7 @@ internal class CgKotlinRenderer(context: CgRendererContext, printer: CgPrinter =
             annotation.accept(this)
         }
 
-        renderClassVisibility(element.id)
+        renderVisibility(element.id)
         renderClassModality(element)
         if (!element.isStatic && element.isNested) {
             print("inner ")
@@ -199,6 +200,12 @@ internal class CgKotlinRenderer(context: CgRendererContext, printer: CgPrinter =
         println(regionEnd)
     }
 
+    override fun visit(element: CgFieldDeclaration) {
+        element.annotation?.accept(this)
+        renderVisibility(element.declaration.variableType)
+        element.declaration.accept(this)
+    }
+
 
     // Property access
 
@@ -580,7 +587,7 @@ internal class CgKotlinRenderer(context: CgRendererContext, printer: CgPrinter =
     override fun escapeNamePossibleKeywordImpl(s: String): String =
         if (isLanguageKeyword(s, context.cgLanguageAssistant)) "`$s`" else s
 
-    override fun renderClassVisibility(classId: ClassId) {
+    override fun renderVisibility(classId: ClassId) {
         when {
             // Kotlin classes are public by default
             classId.isPublic -> Unit

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/renderer/CgVisitor.kt

@@ -76,6 +76,7 @@ import org.utbot.framework.codegen.domain.models.CgSwitchCaseLabel
 import org.utbot.framework.codegen.domain.models.CgClass
 import org.utbot.framework.codegen.domain.models.CgClassBody
 import org.utbot.framework.codegen.domain.models.CgDocRegularLineStmt
+import org.utbot.framework.codegen.domain.models.CgFieldDeclaration
 import org.utbot.framework.codegen.domain.models.CgFormattedString
 import org.utbot.framework.codegen.domain.models.CgFrameworkUtilMethod
 import org.utbot.framework.codegen.domain.models.CgNestedClassesRegion
@@ -185,6 +186,9 @@ interface CgVisitor<R> {
     // Variable declaration
     fun visit(element: CgDeclaration): R
 
+    // Field declaration
+    fun visit(element: CgFieldDeclaration): R
+
     // Variable assignment
     fun visit(element: CgAssignment): R
 

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

@@ -138,6 +138,12 @@ class MockFrameworkManager(context: CgContext) : CgVariableConstructorComponent(
         objectMocker.createMock(model, baseName)
     }
 
+    fun createMockForVariable(model: UtCompositeModel, variable: CgVariable) =
+        withMockFramework {
+            require(model.isMock) { "Mock model is expected in MockObjectConstructor" }
+            objectMocker.mockForVariable(model, variable)
+        }
+
     fun mockNewInstance(mock: UtNewInstanceInstrumentation) {
         staticMocker?.mockNewInstance(mock)
     }
@@ -177,6 +183,12 @@ private class MockitoMocker(context: CgContext) : ObjectMocker(context) {
         val modelClass = getClassOf(model.classId)
         val mockObject = newVar(model.classId, baseName = baseName, isMock = true) { mock(modelClass) }
 
+        mockForVariable(model, mockObject)
+
+        return mockObject
+    }
+
+    fun mockForVariable(model: UtCompositeModel, mockObject: CgVariable) {
         for ((executable, values) in model.mocks) {
             // void method
             if (executable.returnType == voidClassId) {
@@ -201,8 +213,6 @@ private class MockitoMocker(context: CgContext) : ObjectMocker(context) {
                 else -> error("ConstructorId was not expected to appear in simple mocker but got $executable")
             }
         }
-
-        return mockObject
     }
 
     override fun mock(clazz: CgExpression): CgMethodCall =

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

@@ -15,6 +15,7 @@ import org.utbot.framework.codegen.tree.CgMethodConstructor
 import org.utbot.framework.codegen.tree.CgStatementConstructor
 import org.utbot.framework.codegen.tree.CgStatementConstructorImpl
 import org.utbot.framework.codegen.tree.CgVariableConstructor
+import org.utbot.framework.codegen.tree.CgSpringVariableConstructor
 import org.utbot.framework.plugin.api.ClassId
 import org.utbot.framework.plugin.api.CodegenLanguage
 
@@ -44,7 +45,8 @@ abstract class CgLanguageAssistant {
     open fun getCallableAccessManagerBy(context: CgContext): CgCallableAccessManager =
         CgCallableAccessManagerImpl(context)
     open fun getStatementConstructorBy(context: CgContext): CgStatementConstructor = CgStatementConstructorImpl(context)
-    open fun getVariableConstructorBy(context: CgContext): CgVariableConstructor = CgVariableConstructor(context)
+    open fun getVariableConstructorBy(context: CgContext): CgVariableConstructor =
+        if (context.isSpringClass) CgSpringVariableConstructor(context) else CgVariableConstructor(context)
     open fun getMethodConstructorBy(context: CgContext): CgMethodConstructor = CgMethodConstructor(context)
     open fun getCgFieldStateManager(context: CgContext): CgFieldStateManager = CgFieldStateManagerImpl(context)
 

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

@@ -13,6 +13,7 @@ import org.utbot.framework.codegen.domain.models.CgElement
 import org.utbot.framework.codegen.domain.models.CgErrorTestMethod
 import org.utbot.framework.codegen.domain.models.CgExceptionHandler
 import org.utbot.framework.codegen.domain.models.CgExpression
+import org.utbot.framework.codegen.domain.models.CgFieldDeclaration
 import org.utbot.framework.codegen.domain.models.CgForEachLoop
 import org.utbot.framework.codegen.domain.models.CgForLoop
 import org.utbot.framework.codegen.domain.models.CgLoop
@@ -66,7 +67,7 @@ class CgClassBodyBuilder(val classId: ClassId) : CgBuilder<CgClassBody> {
     val methodRegions: MutableList<CgMethodsCluster> = mutableListOf()
     val staticDeclarationRegions: MutableList<CgStaticsRegion> = mutableListOf()
     val nestedClassRegions: MutableList<CgNestedClassesRegion<*>> = mutableListOf()
-    val fields: MutableList<CgDeclaration> = mutableListOf()
+    val fields: MutableList<CgFieldDeclaration> = mutableListOf()
 
     override fun build() = CgClassBody(classId, methodRegions, staticDeclarationRegions, nestedClassRegions, fields)
 }