Save already created mock model in models to variables cache #2494 #2595 (#2532)

Author: EgorkaKulikov

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

@@ -20,22 +20,19 @@ import org.utbot.framework.codegen.domain.builtin.UtilClassFileMethodProvider
 import org.utbot.framework.codegen.domain.builtin.UtilMethodProvider
 import org.utbot.framework.codegen.domain.models.*
 import org.utbot.framework.codegen.services.access.Block
-import org.utbot.framework.codegen.services.access.CgFieldStateManager
 import org.utbot.framework.codegen.tree.EnvironmentFieldStateCache
 import org.utbot.framework.codegen.tree.importIfNeeded
 import org.utbot.framework.plugin.api.BuiltinClassId
 import org.utbot.framework.plugin.api.ClassId
 import org.utbot.framework.codegen.services.language.CgLanguageAssistant
 import org.utbot.framework.codegen.tree.fieldmanager.CgAbstractClassFieldManager
-import org.utbot.framework.codegen.tree.fieldmanager.CgClassFieldManager
 import org.utbot.framework.plugin.api.CodegenLanguage
 import org.utbot.framework.plugin.api.ExecutableId
 import org.utbot.framework.plugin.api.FieldId
 import org.utbot.framework.plugin.api.MethodId
 import org.utbot.framework.plugin.api.MockFramework
 import org.utbot.framework.plugin.api.UtExecution
 import org.utbot.framework.plugin.api.UtModel
-import org.utbot.framework.plugin.api.UtReferenceModel
 import org.utbot.framework.plugin.api.util.id
 import org.utbot.framework.plugin.api.util.isCheckedException
 import org.utbot.framework.plugin.api.util.isSubtypeOf
@@ -344,14 +341,9 @@ interface CgContextOwner {
         return result
     }
 
-    fun updateVariableScope(variable: CgVariable, model: UtModel? = null) {
+    fun rememberVariableForModel(variable: CgVariable, model: UtModel? = null) {
         model?.let {
             valueByUtModelWrapper[it.wrap()] = variable
-            (model as UtReferenceModel).let { refModel ->
-                refModel.id.let {
-                    valueByUtModelWrapper[model.wrap()] = variable
-                }
-            }
         }
     }
 

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

@@ -185,7 +185,8 @@ private class MockitoMocker(context: CgContext) : ObjectMocker(context) {
 
     override fun createMock(model: UtCompositeModel, baseName: String): CgVariable {
         val modelClass = getClassOf(model.classId)
-        val mockObject = newVar(model.classId, baseName = baseName, isMock = true) { mock(modelClass) }
+        val mockObject =
+            newVar(model.classId, model = model, baseName = baseName, isMock = true) { mock(modelClass) }
 
         mockForVariable(model, mockObject)
 

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

@@ -209,7 +209,7 @@ internal class CgStatementConstructorImpl(context: CgContext) :
         }
 
         classRef?.let { declaredClassRefs = declaredClassRefs.put(it, declaration.variable) }
-        updateVariableScope(declaration.variable, model)
+        rememberVariableForModel(declaration.variable, model)
 
         return Either.left(declaration)
     }
@@ -469,7 +469,7 @@ internal class CgStatementConstructorImpl(context: CgContext) :
 
     override fun declareVariable(type: ClassId, name: String): CgVariable =
         CgVariable(name, type).also {
-            updateVariableScope(it)
+            rememberVariableForModel(it)
         }
 
     override fun wrapTypeIfRequired(baseType: ClassId): ClassId =

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

@@ -540,7 +540,7 @@ open class CgVariableConstructor(val context: CgContext) :
     ): Pair<CgVariable, CgDeclaration> {
         val declaration = CgDeclaration(variableType, baseVariableName.toVarName(), initializer.resolve())
         val variable = declaration.variable
-        updateVariableScope(variable)
+        rememberVariableForModel(variable)
         return variable to declaration
     }
 

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

@@ -94,7 +94,7 @@ class PythonCgStatementConstructorImpl(context: CgContext) :
             variableName = name
             initializer = expr
         }
-        updateVariableScope(declaration.variable, model)
+        rememberVariableForModel(declaration.variable, model)
         return Either.left(declaration)
     }
 
@@ -280,7 +280,7 @@ class PythonCgStatementConstructorImpl(context: CgContext) :
 
     override fun declareVariable(type: ClassId, name: String): CgVariable =
         CgVariable(name, type).also {
-            updateVariableScope(it)
+            rememberVariableForModel(it)
         }
 
     override fun guardExpression(baseType: ClassId, expression: CgExpression): ExpressionWithType {