Add CgCustomAssertConstructor to CgComponents

Author: IlyaMuravjov

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

@@ -11,7 +11,9 @@ import org.utbot.framework.codegen.services.access.CgCallableAccessManager
 import org.utbot.framework.codegen.services.access.CgCallableAccessManagerImpl
 import org.utbot.framework.codegen.services.access.CgFieldStateManager
 import org.utbot.framework.codegen.services.access.CgFieldStateManagerImpl
+import org.utbot.framework.codegen.tree.CgCustomAssertConstructor
 import org.utbot.framework.codegen.tree.CgMethodConstructor
+import org.utbot.framework.codegen.tree.CgSimpleCustomAssertConstructor
 import org.utbot.framework.codegen.tree.CgStatementConstructor
 import org.utbot.framework.codegen.tree.CgStatementConstructorImpl
 import org.utbot.framework.codegen.tree.CgVariableConstructor
@@ -47,6 +49,9 @@ interface CgLanguageAssistant {
     fun getVariableConstructorBy(context: CgContext): CgVariableConstructor
 
     fun getMethodConstructorBy(context: CgContext): CgMethodConstructor
+
+    fun getCustomAssertConstructorBy(context: CgContext): CgCustomAssertConstructor
+
     fun getCgFieldStateManager(context: CgContext): CgFieldStateManager
 
     fun getLanguageTestFrameworkManager(): LanguageTestFrameworkManager
@@ -65,4 +70,7 @@ abstract class AbstractCgLanguageAssistant : CgLanguageAssistant {
 
     override fun getMethodConstructorBy(context: CgContext): CgMethodConstructor = CgMethodConstructor(context)
     override fun getCgFieldStateManager(context: CgContext): CgFieldStateManager = CgFieldStateManagerImpl(context)
+
+    override fun getCustomAssertConstructorBy(context: CgContext): CgCustomAssertConstructor =
+        CgSimpleCustomAssertConstructor(context)
 }

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

@@ -22,6 +22,7 @@ object CgComponents {
         mockFrameworkManagers.clear()
         variableConstructors.clear()
         methodConstructors.clear()
+        customAssertConstructors.clear()
     }
 
     private val nameGenerators: IdentityHashMap<CgContext, CgNameGenerator> = IdentityHashMap()
@@ -33,6 +34,7 @@ object CgComponents {
     private val statementConstructors: IdentityHashMap<CgContext, CgStatementConstructor> = IdentityHashMap()
     private val variableConstructors: IdentityHashMap<CgContext, CgVariableConstructor> = IdentityHashMap()
     private val methodConstructors: IdentityHashMap<CgContext, CgMethodConstructor> = IdentityHashMap()
+    private val customAssertConstructors: IdentityHashMap<CgContext, CgCustomAssertConstructor> = IdentityHashMap()
 
     fun getNameGeneratorBy(context: CgContext): CgNameGenerator = nameGenerators.getOrPut(context) {
         context.cgLanguageAssistant.getNameGeneratorBy(context)
@@ -62,4 +64,8 @@ object CgComponents {
     fun getMethodConstructorBy(context: CgContext): CgMethodConstructor = methodConstructors.getOrPut(context) {
         context.cgLanguageAssistant.getMethodConstructorBy(context)
     }
+
+    fun getCustomAssertConstructorBy(context: CgContext): CgCustomAssertConstructor = customAssertConstructors.getOrPut(context) {
+        context.cgLanguageAssistant.getCustomAssertConstructorBy(context)
+    }
 }

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

@@ -0,0 +1,19 @@
+package org.utbot.framework.codegen.tree
+
+import org.utbot.framework.codegen.domain.context.CgContext
+import org.utbot.framework.codegen.domain.context.CgContextOwner
+import org.utbot.framework.codegen.domain.models.CgVariable
+import org.utbot.framework.plugin.api.UtCustomModel
+
+interface CgCustomAssertConstructor {
+    val context: CgContext
+    fun tryConstructCustomAssert(expected: UtCustomModel, actual: CgVariable): Boolean
+}
+
+class CgSimpleCustomAssertConstructor(
+    override val context: CgContext
+) : CgCustomAssertConstructor,
+    CgContextOwner by context {
+    override fun tryConstructCustomAssert(expected: UtCustomModel, actual: CgVariable): Boolean =
+        false
+}

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

@@ -57,6 +57,7 @@ import org.utbot.framework.codegen.services.access.CgCallableAccessManager
 import org.utbot.framework.codegen.services.access.CgFieldStateManagerImpl
 import org.utbot.framework.codegen.services.framework.TestFrameworkManager
 import org.utbot.framework.codegen.tree.CgComponents.getCallableAccessManagerBy
+import org.utbot.framework.codegen.tree.CgComponents.getCustomAssertConstructorBy
 import org.utbot.framework.codegen.tree.CgComponents.getMockFrameworkManagerBy
 import org.utbot.framework.codegen.tree.CgComponents.getNameGeneratorBy
 import org.utbot.framework.codegen.tree.CgComponents.getStatementConstructorBy
@@ -175,6 +176,7 @@ open class CgMethodConstructor(val context: CgContext) : CgContextOwner by conte
     protected val testFrameworkManager = getTestFrameworkManagerBy(context)
 
     protected val variableConstructor = getVariableConstructorBy(context)
+    private val customAssertConstructor = getCustomAssertConstructorBy(context)
     private val mockFrameworkManager = getMockFrameworkManagerBy(context)
 
     private val floatDelta: Float = 1e-6f
@@ -323,11 +325,7 @@ open class CgMethodConstructor(val context: CgContext) : CgContextOwner by conte
                             +thisInstance[executable](*methodArguments.toTypedArray())
                         } else {
                             this.resultModel = resultModel
-
-                            // TODO support custom way of rendering asserts when `resultModel` is `UtCustomModel`
-                            val expected = variableConstructor.getOrCreateVariable(resultModel, "expected")
-                            emptyLineIfNeeded()
-                            assertEquality(expected, actual)
+                            assertEquality(resultModel, actual, emptyLineIfNeeded = true)
                         }
                     }
                     .onFailure { exception -> processExecutionFailure(exception, executionResult) }
@@ -573,9 +571,11 @@ open class CgMethodConstructor(val context: CgContext) : CgContextOwner by conte
         }
 
         if (afterModel !is UtReferenceModel) {
-            val expectedAfter =
-                variableConstructor.getOrCreateVariable(afterModel, "expected" + afterVariable.name.capitalize())
-            assertEquality(expectedAfter, afterVariable)
+            assertEquality(
+                expected = afterModel,
+                actual = afterVariable,
+                expectedVariableName = "expected" + afterVariable.name.capitalize()
+            )
         } else {
             if (beforeVariable != null)
                 testFrameworkManager.assertBoolean(false, beforeVariable equalTo afterVariable)
@@ -1195,6 +1195,19 @@ open class CgMethodConstructor(val context: CgContext) : CgContextOwner by conte
         return ClassIdArrayInfo(classId, nestedElementClassId, dimensions)
     }
 
+    fun assertEquality(
+        expected: UtModel,
+        actual: CgVariable,
+        expectedVariableName: String = "expected",
+        emptyLineIfNeeded: Boolean = false,
+    ) {
+        if (expected !is UtCustomModel || !customAssertConstructor.tryConstructCustomAssert(expected, actual)) {
+            val expectedVariable = variableConstructor.getOrCreateVariable(expected, expectedVariableName)
+            if (emptyLineIfNeeded) emptyLineIfNeeded()
+            assertEquality(expectedVariable, actual)
+        }
+    }
+
     open fun assertEquality(expected: CgValue, actual: CgVariable) {
         when {
             expected.type.isArray -> {

utbot-js/src/main/kotlin/framework/codegen/model/constructor/tree/JsCgMethodConstructor.kt

@@ -69,8 +69,7 @@ class JsCgMethodConstructor(ctx: CgContext) : CgMethodConstructor(ctx) {
                     +thisInstance[method](*methodArguments.toTypedArray())
                 } else {
                     resultModel = result
-                    val expected = variableConstructor.getOrCreateVariable(result, "expected")
-                    assertEquality(expected, actual)
+                    assertEquality(result, actual)
                 }
             }
             .onFailure { exception ->

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

@@ -190,17 +190,21 @@ class PythonCgMethodConstructor(context: CgContext) : CgMethodConstructor(contex
             emptyLineIfNeeded()
         }
         stateAssertions.forEach { (index, it) ->
-            val name = paramNames[executableId]?.get(index) + "_modified"
-            val modifiedArgument = variableConstructor.getOrCreateVariable(it.second, name)
-            assertEquality(modifiedArgument, it.first)
+            assertEquality(
+                expected = it.second,
+                actual = it.first,
+                expectedVariableName = paramNames[executableId]?.get(index) + "_modified"
+            )
         }
         if (assertThisObject.isNotEmpty()) {
             emptyLineIfNeeded()
         }
         assertThisObject.forEach {
-            val name = it.first.name + "_modified"
-            val modifiedThisObject = variableConstructor.getOrCreateVariable(it.second, name)
-            assertEquality(modifiedThisObject, it.first)
+            assertEquality(
+                expected = it.second,
+                actual = it.first,
+                expectedVariableName = it.first.name + "_modified"
+            )
         }
     }