@@ -9,14 +9,15 @@ import org.utbot.framework.codegen.domain.StaticsMocking
99import org.utbot.framework.codegen.domain.TestFramework
1010import org.utbot.framework.codegen.domain.models.CgMethodTestSet
1111import org.utbot.framework.codegen.domain.context.CgContext
12- import org.utbot.framework.codegen.domain.models.SimpleTestClassModelBuilder
13- import org.utbot.framework.codegen.domain.models.SpringTestClassModelBuilder
14- import org.utbot.framework.codegen.domain.models.TestClassModel
12+ import org.utbot.framework.codegen.domain.models.CgClassFile
13+ import org.utbot.framework.codegen.domain.models.builders.SimpleTestClassModelBuilder
14+ import org.utbot.framework.codegen.domain.models.builders.SpringTestClassModelBuilder
1515import org.utbot.framework.codegen.renderer.CgAbstractRenderer
1616import org.utbot.framework.codegen.reports.TestsGenerationReport
1717import org.utbot.framework.codegen.tree.CgSimpleTestClassConstructor
1818import org.utbot.framework.codegen.tree.ututils.UtilClassKind
1919import org.utbot.framework.codegen.services.language.CgLanguageAssistant
20+ import org.utbot.framework.codegen.tree.CgSpringTestClassConstructor
2021import org.utbot.framework.plugin.api.ClassId
2122import org.utbot.framework.plugin.api.CodegenLanguage
2223import org.utbot.framework.plugin.api.ExecutableId
@@ -75,33 +76,60 @@ open class CodeGenerator(
7576 val cgTestSets = testSets.map { CgMethodTestSet (it) }.toList()
7677 return withCustomContext(testClassCustomName) {
7778 context.withTestClassFileScope {
78- val astConstructor = CgSimpleTestClassConstructor (context)
79- val renderer = CgAbstractRenderer .makeRenderer(context)
79+ if (context.isSpringClass) {
80+ generateForSpringClass(cgTestSets)
81+ } else {
82+ generateForSimpleClass(cgTestSets)
83+ }
84+ }
85+ }
86+ }
8087
81- val testClassModelBuilder =
82- if (context.isSpringClass) SpringTestClassModelBuilder () else SimpleTestClassModelBuilder ()
88+ private fun generateForSimpleClass (testSets : List <CgMethodTestSet >): CodeGeneratorResult {
89+ val astConstructor = CgSimpleTestClassConstructor (context)
90+ val testClassModel = SimpleTestClassModelBuilder ().createTestClassModel(classUnderTest, testSets)
8391
84- val testClassModel = testClassModelBuilder.createClassModel(classUnderTest, cgTestSets)
92+ logger.info { " Code generation phase started at ${now()} " }
93+ val testClassFile = astConstructor.construct(testClassModel)
94+ logger.info { " Code generation phase finished at ${now()} " }
8595
86- fun now () = LocalDateTime .now().format( DateTimeFormatter .ofPattern( " HH:mm:ss.SSS " ) )
96+ val generatedCode = renderToString(testClassFile )
8797
88- logger.info { " Code generation phase started at ${now()} " }
89- val testClassFile = astConstructor.construct(testClassModel)
90- logger.info { " Code generation phase finished at ${now()} " }
98+ return CodeGeneratorResult (
99+ generatedCode = generatedCode,
100+ utilClassKind = UtilClassKind .fromCgContextOrNull(context),
101+ testsGenerationReport = astConstructor.testsGenerationReport
102+ )
103+ }
91104
92- logger.info { " Rendering phase started at ${now()} " }
93- testClassFile.accept(renderer )
94- logger.info { " Rendering phase finished at ${now()} " }
105+ private fun generateForSpringClass ( testSets : List < CgMethodTestSet >): CodeGeneratorResult {
106+ val astConstructor = CgSpringTestClassConstructor (context )
107+ val testClassModel = SpringTestClassModelBuilder ().createTestClassModel(classUnderTest, testSets)
95108
96- CodeGeneratorResult (
97- generatedCode = renderer.toString(),
98- utilClassKind = UtilClassKind .fromCgContextOrNull(context),
99- testsGenerationReport = astConstructor.testsGenerationReport
100- )
101- }
102- }
109+ logger.info { " Code generation phase started at ${now()} " }
110+ val testClassFile = astConstructor.construct(testClassModel)
111+ logger.info { " Code generation phase finished at ${now()} " }
112+
113+ val generatedCode = renderToString(testClassFile)
114+
115+ return CodeGeneratorResult (
116+ generatedCode = generatedCode,
117+ utilClassKind = UtilClassKind .fromCgContextOrNull(context),
118+ testsGenerationReport = TestsGenerationReport ()
119+ )
120+ }
121+
122+ private fun renderToString (testClassFile : CgClassFile ): String {
123+ logger.info { " Rendering phase started at ${now()} " }
124+ val renderer = CgAbstractRenderer .makeRenderer(context)
125+ testClassFile.accept(renderer)
126+ logger.info { " Rendering phase finished at ${now()} " }
127+
128+ return renderer.toString()
103129 }
104130
131+ private fun now () = LocalDateTime .now().format(DateTimeFormatter .ofPattern(" HH:mm:ss.SSS" ))
132+
105133 /* *
106134 * Wrapper function that configures context as needed for utbot-online:
107135 * - turns on imports optimization in code generator
0 commit comments