Store kind of util class in a comment to obtain the kind of an existing util class

Author: ArsenHD

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

@@ -82,8 +82,7 @@ class CodeGenerator(
             CodeGeneratorResult(
                 generatedCode = renderClassFile(testClassFile),
                 utilClassKind = UtilClassKind.fromCgContextOrNull(context),
-                testsGenerationReport = testClassFile.testsGenerationReport,
-                mockFrameworkUsed = context.mockFrameworkUsed
+                testsGenerationReport = testClassFile.testsGenerationReport
             )
         }
     }
@@ -117,14 +116,12 @@ class CodeGenerator(
  * @property generatedCode the source code of the test class
  * @property utilClassKind the kind of util class if it is required, otherwise - null
  * @property testsGenerationReport some info about test generation process
- * @property mockFrameworkUsed flag indicating whether any mock objects have been created during code generation ot not
  */
 data class CodeGeneratorResult(
     val generatedCode: String,
     // null if no util class needed, e.g. when we are generating utils directly into test class
     val utilClassKind: UtilClassKind?,
     val testsGenerationReport: TestsGenerationReport,
-    val mockFrameworkUsed: Boolean = false
 )
 
 /**
@@ -162,15 +159,40 @@ sealed class UtilClassKind(
     val utilClassVersionComment: CgComment
         get() = CgSingleLineComment("$UTIL_CLASS_VERSION_COMMENT_PREFIX${utilClassVersion}")
 
+
+    /**
+     * The comment specifying the kind of util class being generated.
+     *
+     * @see utilClassKindCommentText
+     */
+    val utilClassKindComment: CgComment
+        get() = CgSingleLineComment(utilClassKindCommentText)
+
+    /**
+     * The text of comment specifying the kind of util class.
+     * At the moment, there are two kinds: [RegularUtUtils] (without Mockito) and [UtUtilsWithMockito].
+     *
+     * This comment is needed when the plugin decides whether to overwrite an existing util class or not.
+     * When making that decision, it is important to determine if the existing class uses mocks or not,
+     * and this comment will help do that.
+     */
+    abstract val utilClassKindCommentText: String
+
     /**
      * A kind of regular UtUtils class. "Regular" here means that this class does not use a mock framework.
      */
-    object RegularUtUtils : UtilClassKind(UtilClassFileMethodProvider, mockFrameworkUsed = false, priority = 0)
+    object RegularUtUtils : UtilClassKind(UtilClassFileMethodProvider, mockFrameworkUsed = false, priority = 0) {
+        override val utilClassKindCommentText: String
+            get() = "This is a regular UtUtils class (without mock framework usage)"
+    }
 
     /**
      * A kind of UtUtils class that uses a mock framework. At the moment the framework is Mockito.
      */
-    object UtUtilsWithMockito : UtilClassKind(UtilClassFileMethodProvider, mockFrameworkUsed = true, priority = 1)
+    object UtUtilsWithMockito : UtilClassKind(UtilClassFileMethodProvider, mockFrameworkUsed = true, priority = 1) {
+        override val utilClassKindCommentText: String
+            get() = "This is UtUtils class with Mockito support"
+    }
 
     override fun compareTo(other: UtilClassKind): Int {
         return priority.compareTo(other.priority)
@@ -196,6 +218,14 @@ sealed class UtilClassKind(
          */
         const val UTIL_CLASS_VERSION_COMMENT_PREFIX = "UtUtils class version: "
 
+        fun utilClassKindByCommentOrNull(comment: String): UtilClassKind? {
+            return when (comment) {
+                RegularUtUtils.utilClassKindCommentText -> RegularUtUtils
+                UtUtilsWithMockito.utilClassKindCommentText -> UtUtilsWithMockito
+                else -> null
+            }
+        }
+
         /**
          * Check if an util class is required, and if so, what kind.
          * @return `null` if [CgContext.utilMethodProvider] is not [UtilClassFileMethodProvider],

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/model/constructor/tree/CgUtilClassConstructor.kt

@@ -23,6 +23,7 @@ internal object CgUtilClassConstructor {
                 id = utUtilsClassId
                 body = buildRegularClassBody {
                     content += utilClassKind.utilClassVersionComment
+                    content += utilClassKind.utilClassKindComment
                     content += utilMethodProvider.utilMethodIds.map { CgUtilMethod(it) }
                 }
             }

utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/generator/CodeGenerationController.kt

@@ -96,19 +96,9 @@ object CodeGenerationController {
 
     private class UtilClassListener {
         var requiredUtilClassKind: UtilClassKind? = null
-        var mockFrameworkUsed: Boolean = false
 
         fun onTestClassGenerated(result: CodeGeneratorResult) {
             requiredUtilClassKind = maxOfNullable(requiredUtilClassKind, result.utilClassKind)
-            mockFrameworkUsed = maxOf(mockFrameworkUsed, result.mockFrameworkUsed)
-        }
-
-        private fun <T : Comparable<T>> maxOfNullable(a: T?, b: T?): T? {
-            return when {
-                a == null -> b
-                b == null -> a
-                else -> maxOf(a, b)
-            }
         }
     }
 
@@ -153,12 +143,12 @@ object CodeGenerationController {
 
         run(EDT_LATER) {
             waitForCountDown(latch, timeout = 100, timeUnit = TimeUnit.MILLISECONDS) {
-                val mockFrameworkUsed = utilClassListener.mockFrameworkUsed
-                val utilClassKind = utilClassListener.requiredUtilClassKind
+                val requiredUtilClassKind = utilClassListener.requiredUtilClassKind
                     ?: return@waitForCountDown // no util class needed
 
                 val existingUtilClass = model.codegenLanguage.getUtilClassOrNull(model.project, model.testModule)
-                if (shouldCreateOrUpdateUtilClass(existingUtilClass, mockFrameworkUsed, utilClassKind)) {
+                val utilClassKind = newUtilClassKindOrNull(existingUtilClass, requiredUtilClassKind)
+                if (utilClassKind != null) {
                     createOrUpdateUtilClass(
                         testDirectory = baseTestDirectory,
                         utilClassKind = utilClassKind,
@@ -196,40 +186,50 @@ object CodeGenerationController {
         }
     }
 
-    private fun shouldCreateOrUpdateUtilClass(
-        existingUtilClass: PsiFile?,
-        mockFrameworkUsed: Boolean,
-        requiredUtilClassKind: UtilClassKind
-    ): Boolean {
-        val mockFrameworkNotUsed = !mockFrameworkUsed
-
-        val utilClassExists = existingUtilClass != null
-
-        if (!utilClassExists) {
-            // If no util class exists, then we should create a new one.
-            return true
+    /**
+     * This method decides whether to overwrite an existing util class with a new one. And if so, then with what kind of util class.
+     * - If no util class exists, then we generate a new one.
+     * - If existing util class' version is out of date, then we overwrite it with a new one.
+     * But we use the maximum of two kinds (existing and the new one) to avoid problems with mocks.
+     * - If existing util class is up-to-date **and** has a greater or equal priority than the new one,
+     * then we do not need to overwrite it (return null).
+     * - Lastly, if the new util class kind has a greater priority than the existing one,
+     * then we do overwrite it with a newer version.
+     *
+     * @param existingUtilClass a [PsiFile] representing a file of an existing util class. If it does not exist, then [existingUtilClass] is `null`.
+     * @param requiredUtilClassKind the kind of the new util class that we attempt to generate.
+     * @return an [UtilClassKind] of a new util class that will be created or `null`, if no new util class is needed.
+     */
+    private fun newUtilClassKindOrNull(existingUtilClass: PsiFile?, requiredUtilClassKind: UtilClassKind): UtilClassKind? {
+        if (existingUtilClass == null) {
+            // If no util class exists, then we should create a new one with the given kind.
+            return requiredUtilClassKind
         }
 
-        val existingUtilClassVersion = existingUtilClass?.utilClassVersionOrNull
+        val existingUtilClassVersion = existingUtilClass.utilClassVersionOrNull ?: return requiredUtilClassKind
         val newUtilClassVersion = requiredUtilClassKind.utilClassVersion
         val versionIsUpdated = existingUtilClassVersion != newUtilClassVersion
 
+        val existingUtilClassKind = existingUtilClass.utilClassKindOrNull ?: return requiredUtilClassKind
+
         if (versionIsUpdated) {
-            // If an existing util class is out of date,
-            // then we must overwrite it with a newer version.
-            return true
+            // If an existing util class is out of date, then we must overwrite it with a newer version.
+            // But we choose the kind with more priority, because it is possible that
+            // the existing util class needed mocks, but the new one doesn't.
+            // In this case we still want to support mocks, because the previously generated tests
+            // expect that the util class does support them.
+            return maxOfNullable(existingUtilClassKind, requiredUtilClassKind)
         }
 
-        if (mockFrameworkNotUsed) {
-            // If util class already exists and mock framework is not used,
-            // then existing util class is enough, and we don't need to generate a new one.
-            // That's because both regular and mock versions of util class can work
-            // with tests that do not use mocks, so we do not have to worry about
-            // version of util class that we have at the moment.
-            return false
+        if (requiredUtilClassKind <= existingUtilClassKind) {
+            // If the existing util class kind has a greater or equal priority than the new one we attempt to generate,
+            // then we should not do anything. The existing util class is already enough.
+            return null
         }
 
-        return true
+        // The last case. The existing util class has a strictly less priority than the new one.
+        // So we generate the new one to overwrite the previous one with it.
+        return requiredUtilClassKind
     }
 
     /**
@@ -337,7 +337,7 @@ object CodeGenerationController {
     }
 
     /**
-     * Util class must have a comment that specifies the version of UTBot it was generated with.
+     * Util class must have a comment that specifies its version.
      * This property represents the version specified by this comment if it exists. Otherwise, the property is `null`.
      */
     private val PsiFile.utilClassVersionOrNull: String?
@@ -354,6 +354,23 @@ object CodeGenerationController {
                 ?.trim()
         }
 
+    /**
+     * Util class must have a comment that specifies its kind.
+     * This property obtains the kind specified by this comment if it exists. Otherwise, the property is `null`.
+     */
+    private val PsiFile.utilClassKindOrNull: UtilClassKind?
+        get() = runReadAction {
+            val utilClass = (this as? PsiClassOwner)
+                ?.classes
+                ?.firstOrNull()
+                ?: return@runReadAction null
+
+            utilClass.childrenOfType<PsiComment>()
+                .map { comment -> comment.text }
+                .mapNotNull { text -> UtilClassKind.utilClassKindByCommentOrNull(text) }
+                .firstOrNull()
+        }
+
     /**
      * @param srcClass class under test
      * @return name of the package of a given [srcClass].
@@ -878,4 +895,12 @@ object CodeGenerationController {
             title = "Failed to Create Class"
         )
     }
+
+    private fun <T : Comparable<T>> maxOfNullable(a: T?, b: T?): T? {
+        return when {
+            a == null -> b
+            b == null -> a
+            else -> maxOf(a, b)
+        }
+    }
 }