Apply review fixes

utbot-framework/src/main/kotlin/org/utbot/framework/assemble/AssembleModelGenerator.kt

@@ -343,23 +343,23 @@ class AssembleModelGenerator(private val basePackageName: String) {
     private fun assembleMockCompositeModel(compositeModel: UtCompositeModel): UtCompositeModel {
         // We have to create a model before the construction of the fields to avoid
         // infinite recursion when some mock contains itself as a field.
-        val assembledModel = UtCompositeModel(
+        val assembledCompositeModel = UtCompositeModel(
             compositeModel.id,
             compositeModel.classId,
             isMock = true,
         )
 
-        instantiatedModels[compositeModel] = assembledModel
+        instantiatedModels[compositeModel] = assembledCompositeModel
 
         val fields = compositeModel.fields.mapValues { assembleModel(it.value) }.toMutableMap()
         val mockBehaviour = compositeModel.mocks
             .mapValues { models -> models.value.map { assembleModel(it) } }
             .toMutableMap()
 
-        assembledModel.fields += fields
-        assembledModel.mocks += mockBehaviour
+        assembledCompositeModel.fields += fields
+        assembledCompositeModel.mocks += mockBehaviour
 
-        return assembledModel
+        return assembledCompositeModel
     }
 
     /**

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

@@ -9,12 +9,15 @@ import org.utbot.framework.codegen.domain.StaticsMocking
 import org.utbot.framework.codegen.domain.TestFramework
 import org.utbot.framework.codegen.domain.models.CgMethodTestSet
 import org.utbot.framework.codegen.domain.context.CgContext
-import org.utbot.framework.codegen.domain.models.TestClassModelBuilder
+import org.utbot.framework.codegen.domain.models.CgClassFile
+import org.utbot.framework.codegen.domain.models.builders.SimpleTestClassModelBuilder
+import org.utbot.framework.codegen.domain.models.builders.SpringTestClassModelBuilder
 import org.utbot.framework.codegen.renderer.CgAbstractRenderer
 import org.utbot.framework.codegen.reports.TestsGenerationReport
 import org.utbot.framework.codegen.tree.CgSimpleTestClassConstructor
 import org.utbot.framework.codegen.tree.ututils.UtilClassKind
 import org.utbot.framework.codegen.services.language.CgLanguageAssistant
+import org.utbot.framework.codegen.tree.CgSpringTestClassConstructor
 import org.utbot.framework.plugin.api.ClassId
 import org.utbot.framework.plugin.api.CodegenLanguage
 import org.utbot.framework.plugin.api.ExecutableId
@@ -73,31 +76,60 @@ open class CodeGenerator(
         val cgTestSets = testSets.map { CgMethodTestSet(it) }.toList()
         return withCustomContext(testClassCustomName) {
             context.withTestClassFileScope {
-                val astConstructor = CgSimpleTestClassConstructor(context)
-                val renderer = CgAbstractRenderer.makeRenderer(context)
-                val testClassModelBuilder = TestClassModelBuilder(context.isSpringClass)
+                if (context.isSpringClass) {
+                    generateForSpringClass(cgTestSets)
+                } else {
+                    generateForSimpleClass(cgTestSets)
+                }
+            }
+        }
+    }
 
-                val testClassModel = testClassModelBuilder.createClassModel(classUnderTest, cgTestSets)
+    private fun generateForSimpleClass(testSets: List<CgMethodTestSet>): CodeGeneratorResult {
+        val astConstructor = CgSimpleTestClassConstructor(context)
+        val testClassModel = SimpleTestClassModelBuilder().createTestClassModel(classUnderTest, testSets)
 
-                fun now() = LocalDateTime.now().format(DateTimeFormatter.ofPattern("HH:mm:ss.SSS"))
+        logger.info { "Code generation phase started at ${now()}" }
+        val testClassFile = astConstructor.construct(testClassModel)
+        logger.info { "Code generation phase finished at ${now()}" }
 
-                logger.info { "Code generation phase started at ${now()}" }
-                val testClassFile = astConstructor.construct(testClassModel)
-                logger.info { "Code generation phase finished at ${now()}" }
+        val generatedCode = renderToString(testClassFile)
 
-                logger.info { "Rendering phase started at ${now()}" }
-                testClassFile.accept(renderer)
-                logger.info { "Rendering phase finished at ${now()}" }
+        return CodeGeneratorResult(
+            generatedCode = generatedCode,
+            utilClassKind = UtilClassKind.fromCgContextOrNull(context),
+            testsGenerationReport = astConstructor.testsGenerationReport
+        )
+    }
 
-                CodeGeneratorResult(
-                    generatedCode = renderer.toString(),
-                    utilClassKind = UtilClassKind.fromCgContextOrNull(context),
-                    testsGenerationReport = astConstructor.testsGenerationReport
-                )
-            }
-        }
+    private fun generateForSpringClass(testSets: List<CgMethodTestSet>): CodeGeneratorResult {
+        val astConstructor = CgSpringTestClassConstructor(context)
+        val testClassModel = SpringTestClassModelBuilder().createTestClassModel(classUnderTest, testSets)
+
+        logger.info { "Code generation phase started at ${now()}" }
+        val testClassFile = astConstructor.construct(testClassModel)
+        logger.info { "Code generation phase finished at ${now()}" }
+
+        val generatedCode = renderToString(testClassFile)
+
+        return CodeGeneratorResult(
+            generatedCode = generatedCode,
+            utilClassKind = UtilClassKind.fromCgContextOrNull(context),
+            testsGenerationReport = TestsGenerationReport()
+        )
+    }
+
+    private fun renderToString(testClassFile: CgClassFile): String {
+        logger.info { "Rendering phase started at ${now()}" }
+        val renderer = CgAbstractRenderer.makeRenderer(context)
+        testClassFile.accept(renderer)
+        logger.info { "Rendering phase finished at ${now()}" }
+
+        return renderer.toString()
     }
 
+    private fun now() = LocalDateTime.now().format(DateTimeFormatter.ofPattern("HH:mm:ss.SSS"))
+
     /**
      * Wrapper function that configures context as needed for utbot-online:
      * - turns on imports optimization in code generator

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

@@ -26,7 +26,7 @@ import org.utbot.framework.codegen.domain.models.CgMethodTestSet
 import org.utbot.framework.codegen.domain.builtin.TestClassUtilMethodProvider
 import org.utbot.framework.codegen.domain.builtin.UtilClassFileMethodProvider
 import org.utbot.framework.codegen.domain.builtin.UtilMethodProvider
-import org.utbot.framework.codegen.domain.models.TestClassModel
+import org.utbot.framework.codegen.domain.models.SimpleTestClassModel
 import org.utbot.framework.codegen.domain.models.CgParameterKind
 import org.utbot.framework.codegen.services.access.Block
 import org.utbot.framework.codegen.tree.EnvironmentFieldStateCache
@@ -295,7 +295,7 @@ interface CgContextOwner {
      * This method does almost all the same as [withTestClassScope], but for nested test classes.
      * The difference is that instead of working with [outerMostTestClassContext] it works with [currentTestClassContext].
      */
-    fun <R> withNestedClassScope(testClassModel: TestClassModel, block: () -> R): R
+    fun <R> withNestedClassScope(testClassModel: SimpleTestClassModel, block: () -> R): R
 
     /**
      * Set [mockFrameworkUsed] flag to true if the block is successfully executed
@@ -434,7 +434,7 @@ interface CgContextOwner {
  */
 data class CgContext(
     override val classUnderTest: ClassId,
-    override val isSpringClass: Boolean = false,
+    override val isSpringClass: Boolean = true,
     val generateUtilClassFile: Boolean = false,
     override var currentExecutable: ExecutableId? = null,
     override val collectedExceptions: MutableSet<ClassId> = mutableSetOf(),
@@ -542,7 +542,7 @@ data class CgContext(
         }
     }
 
-    override fun <R> withNestedClassScope(testClassModel: TestClassModel, block: () -> R): R {
+    override fun <R> withNestedClassScope(testClassModel: SimpleTestClassModel, block: () -> R): R {
         val previousCurrentTestClassInfo = currentTestClassContext
         val previousCurrentTestClass = currentTestClass
         currentTestClass = createClassIdForNestedClass(testClassModel)
@@ -555,7 +555,7 @@ data class CgContext(
         }
     }
 
-    private fun createClassIdForNestedClass(testClassModel: TestClassModel): ClassId {
+    private fun createClassIdForNestedClass(testClassModel: SimpleTestClassModel): ClassId {
         val simpleName = "${testClassModel.classUnderTest.simpleName}Test"
         return BuiltinClassId(
             canonicalName = currentTestClass.canonicalName + "." + simpleName,

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

@@ -2,19 +2,34 @@ package org.utbot.framework.codegen.domain.models
 
 import org.utbot.framework.plugin.api.ClassId
 
-// TODO: seems like this class needs to be renamed
 /**
  * Stores method test sets in a structure that replicates structure of their methods in [classUnderTest].
  * I.e., if some method is declared in nested class of [classUnderTest], its testset will be put
  * in [TestClassModel] in one of [nestedClasses]
+ */
+abstract class TestClassModel(
+    val classUnderTest: ClassId,
+    val methodTestSets: List<CgMethodTestSet>,
+    val nestedClasses: List<SimpleTestClassModel>,
+)
+
+open class SimpleTestClassModel(
+    classUnderTest: ClassId,
+    methodTestSets: List<CgMethodTestSet>,
+    nestedClasses: List<SimpleTestClassModel> = listOf(),
+): TestClassModel(classUnderTest, methodTestSets, nestedClasses)
+
+/**
+ * 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 TestClassModel(
-    val classUnderTest: ClassId,
-    val methodTestSets: List<CgMethodTestSet>,
-    val nestedClasses: List<TestClassModel> = listOf(),
+class SpringTestClassModel(
+    classUnderTest: ClassId,
+    methodTestSets: List<CgMethodTestSet>,
+    nestedClasses: List<SimpleTestClassModel>,
     val injectingMocksClass: ClassId? = null,
     val mockedClasses: Set<ClassId> = setOf(),
-)
+): TestClassModel(classUnderTest, methodTestSets, nestedClasses)
+