CgContext is not a data class!!!

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

@@ -141,10 +141,7 @@ open class CodeGenerator(
     fun <R> withCustomContext(testClassCustomName: String? = null, block: () -> R): R {
         val prevContext = context
         return try {
-            context = prevContext.copy(
-                    shouldOptimizeImports = true,
-                    testClassCustomName = testClassCustomName
-            )
+            context = prevContext.customCopy(shouldOptimizeImports = true, testClassCustomName = testClassCustomName)
             block()
         } finally {
             context = prevContext

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

@@ -445,7 +445,7 @@ interface CgContextOwner {
 /**
  * Context with current code generation info
  */
-data class CgContext(
+class CgContext(
     override val classUnderTest: ClassId,
     override val projectType: ProjectType,
     val generateUtilClassFile: Boolean = false,
@@ -594,7 +594,6 @@ data class CgContext(
         mockFrameworkUsed = false
     }
 
-
     override var valueByModel: IdentityHashMap<UtModel, CgValue> = IdentityHashMap()
 
     override var valueByModelId: MutableMap<ModelId, CgValue> = mutableMapOf()
@@ -605,4 +604,47 @@ data class CgContext(
 
     override val utilMethodsUsed: Boolean
         get() = requiredUtilMethods.isNotEmpty()
+
+    fun customCopy(shouldOptimizeImports: Boolean, testClassCustomName: String?) = CgContext(
+        shouldOptimizeImports = shouldOptimizeImports,
+        testClassCustomName = testClassCustomName,
+
+        classUnderTest = this.classUnderTest,
+        projectType = this.projectType,
+        generateUtilClassFile = this.generateUtilClassFile,
+        currentExecutable = this.currentExecutable,
+        collectedExceptions =this.collectedExceptions,
+        collectedMethodAnnotations = this.collectedMethodAnnotations,
+        collectedImports = this.collectedImports,
+        importedStaticMethods = this.importedStaticMethods,
+        importedClasses = this.importedClasses,
+        requiredUtilMethods = this.requiredUtilMethods,
+        testMethods = this.testMethods,
+        existingMethodNames = this.existingMethodNames,
+        prevStaticFieldValues = this.prevStaticFieldValues,
+        paramNames = this.paramNames,
+        currentExecution = this.currentExecution,
+        testFramework = this.testFramework,
+        mockFramework = this.mockFramework,
+        staticsMocking = this.staticsMocking,
+        forceStaticMocking = this.forceStaticMocking,
+        generateWarningsForStaticMocking = this.generateWarningsForStaticMocking,
+        codegenLanguage = this.codegenLanguage,
+        cgLanguageAssistant = this.cgLanguageAssistant,
+        parametrizedTestSource = this.parametrizedTestSource,
+        mockFrameworkUsed = this.mockFrameworkUsed,
+        currentBlock = this.currentBlock,
+        existingVariableNames = this.existingVariableNames,
+        declaredClassRefs = this.declaredClassRefs,
+        declaredExecutableRefs = this.declaredExecutableRefs,
+        declaredFieldRefs = this.declaredFieldRefs,
+        thisInstance = this.thisInstance,
+        methodArguments = this.methodArguments,
+        codeGenerationErrors = this.codeGenerationErrors,
+        testClassPackageName = this.testClassPackageName,
+        runtimeExceptionTestsBehaviour = this.runtimeExceptionTestsBehaviour,
+        hangingTestsTimeout = this.hangingTestsTimeout,
+        enableTestsTimeout = this.enableTestsTimeout,
+        containsReflectiveCall = this.containsReflectiveCall,
+    )
 }

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

@@ -5,6 +5,7 @@ import org.utbot.framework.codegen.services.CgNameGenerator
 import org.utbot.framework.codegen.services.access.CgCallableAccessManager
 import org.utbot.framework.codegen.services.framework.MockFrameworkManager
 import org.utbot.framework.codegen.services.framework.TestFrameworkManager
+import java.util.*
 
 object CgComponents {
 
@@ -23,33 +24,42 @@ object CgComponents {
         methodConstructors.clear()
     }
 
-    private val nameGenerators: MutableMap<CgContext, CgNameGenerator> = mutableMapOf()
-    private val statementConstructors: MutableMap<CgContext, CgStatementConstructor> = mutableMapOf()
-    private val callableAccessManagers: MutableMap<CgContext, CgCallableAccessManager> = mutableMapOf()
-    private val testFrameworkManagers: MutableMap<CgContext, TestFrameworkManager> = mutableMapOf()
-    private val mockFrameworkManagers: MutableMap<CgContext, MockFrameworkManager> = mutableMapOf()
+    private val nameGenerators: IdentityHashMap<CgContext, CgNameGenerator> = IdentityHashMap()
 
-    private val variableConstructors: MutableMap<CgContext, CgVariableConstructor> = mutableMapOf()
-    private val methodConstructors: MutableMap<CgContext, CgMethodConstructor> = mutableMapOf()
+    private val callableAccessManagers: IdentityHashMap<CgContext, CgCallableAccessManager> = IdentityHashMap()
+    private val testFrameworkManagers: IdentityHashMap<CgContext, TestFrameworkManager> = IdentityHashMap()
+    private val mockFrameworkManagers: IdentityHashMap<CgContext, MockFrameworkManager> = IdentityHashMap()
 
-    fun getNameGeneratorBy(context: CgContext) = nameGenerators.getOrPut(context) {
+    private val statementConstructors: IdentityHashMap<CgContext, CgStatementConstructor> = IdentityHashMap()
+    private val variableConstructors: IdentityHashMap<CgContext, CgVariableConstructor> = IdentityHashMap()
+    private val methodConstructors: IdentityHashMap<CgContext, CgMethodConstructor> = IdentityHashMap()
+
+    fun getNameGeneratorBy(context: CgContext): CgNameGenerator = nameGenerators.getOrPut(context) {
         context.cgLanguageAssistant.getNameGeneratorBy(context)
     }
-    fun getCallableAccessManagerBy(context: CgContext) = callableAccessManagers.getOrPut(context) {
+
+    fun getCallableAccessManagerBy(context: CgContext): CgCallableAccessManager = callableAccessManagers.getOrPut(context) {
         context.cgLanguageAssistant.getCallableAccessManagerBy(context)
     }
-    fun getStatementConstructorBy(context: CgContext) = statementConstructors.getOrPut(context) {
+
+    fun getStatementConstructorBy(context: CgContext): CgStatementConstructor = statementConstructors.getOrPut(context) {
         context.cgLanguageAssistant.getStatementConstructorBy(context)
     }
 
-    fun getTestFrameworkManagerBy(context: CgContext) =
-        testFrameworkManagers.getOrDefault(context, context.cgLanguageAssistant.getLanguageTestFrameworkManager().managerByFramework(context))
+    fun getTestFrameworkManagerBy(context: CgContext): TestFrameworkManager =
+        testFrameworkManagers.getOrDefault(
+            context,
+            context.cgLanguageAssistant.getLanguageTestFrameworkManager().managerByFramework(context)
+        )
 
-    fun getMockFrameworkManagerBy(context: CgContext) = mockFrameworkManagers.getOrPut(context) { MockFrameworkManager(context) }
-    fun getVariableConstructorBy(context: CgContext) = variableConstructors.getOrPut(context) {
+    fun getMockFrameworkManagerBy(context: CgContext): MockFrameworkManager =
+        mockFrameworkManagers.getOrPut(context) { MockFrameworkManager(context) }
+
+    fun getVariableConstructorBy(context: CgContext): CgVariableConstructor = variableConstructors.getOrPut(context) {
         context.cgLanguageAssistant.getVariableConstructorBy(context)
     }
-    fun getMethodConstructorBy(context: CgContext) = methodConstructors.getOrPut(context) {
+
+    fun getMethodConstructorBy(context: CgContext): CgMethodConstructor = methodConstructors.getOrPut(context) {
         context.cgLanguageAssistant.getMethodConstructorBy(context)
     }
 }

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

@@ -1,6 +1,5 @@
 package org.utbot.framework.codegen.tree
 
-import com.jetbrains.rd.util.firstOrNull
 import org.utbot.framework.codegen.domain.context.CgContext
 import org.utbot.framework.codegen.domain.models.CgValue
 import org.utbot.framework.codegen.domain.models.CgVariable
@@ -13,8 +12,6 @@ class CgSpringVariableConstructor(context: CgContext) : CgVariableConstructor(co
     val injectedMocksModelsVariables: MutableMap<Set<UtModel>, CgValue> = mutableMapOf()
     val mockedModelsVariables: MutableMap<Set<UtModel>, CgValue> = mutableMapOf()
 
-    private val mockFrameworkManager = CgComponents.getMockFrameworkManagerBy(context)
-
     override fun getOrCreateVariable(model: UtModel, name: String?): CgValue {
         val alreadyCreatedInjectMocks = findCgValueByModel(model, injectedMocksModelsVariables)
         if (alreadyCreatedInjectMocks != null) {

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

@@ -75,7 +75,7 @@ open class CgVariableConstructor(val context: CgContext) :
     CgStatementConstructor by getStatementConstructorBy(context) {
 
     private val nameGenerator = getNameGeneratorBy(context)
-    private val mockFrameworkManager = getMockFrameworkManagerBy(context)
+    protected val mockFrameworkManager = getMockFrameworkManagerBy(context)
 
     /**
      * Take already created CgValue or construct either a new [CgVariable] or new [CgLiteral] for the given model.

utbot-js/src/main/kotlin/codegen/JsCodeGenerator.kt

@@ -64,10 +64,7 @@ class JsCodeGenerator(
     private fun <R> withCustomContext(testClassCustomName: String? = null, block: () -> R): R {
         val prevContext = context
         return try {
-            context = prevContext.copy(
-                shouldOptimizeImports = true,
-                testClassCustomName = testClassCustomName
-            )
+            context = prevContext.customCopy(shouldOptimizeImports = true, testClassCustomName = testClassCustomName)
             block()
         } finally {
             context = prevContext