Properly handle canHaveRedundantOrMissingMocks in instrumentation

Author: IlyaMuravjov

utbot-framework-api/src/main/kotlin/org/utbot/framework/plugin/api/Api.kt

@@ -509,7 +509,6 @@ data class UtCompositeModel(
     val isMock: Boolean,
     val fields: MutableMap<FieldId, UtModel> = mutableMapOf(),
     val mocks: MutableMap<ExecutableId, List<UtModel>> = mutableMapOf(),
-    // TODO handle it in instrumentation & codegen
     val canHaveRedundantOrMissingMocks: Boolean = false,
 ) : UtReferenceModel(id, classId) {
     //TODO: SAT-891 - rewrite toString() method
@@ -767,13 +766,17 @@ abstract class UtCustomModel(
     classId: ClassId,
     modelName: String = id.toString(),
     override val origin: UtCompositeModel? = null,
-) : UtModelWithCompositeOrigin(id, classId, modelName, origin)
+) : UtModelWithCompositeOrigin(id, classId, modelName, origin) {
+    abstract val dependencies: Collection<UtModel>
+}
 
 object UtSpringContextModel : UtCustomModel(
     id = null,
     classId = SpringModelUtils.applicationContextClassId,
     modelName = "applicationContext"
 ) {
+    override val dependencies: Collection<UtModel> get() = emptySet()
+
     // NOTE that overriding equals is required just because without it
     // we will lose equality for objects after deserialization
     override fun equals(other: Any?): Boolean = other is UtSpringContextModel
@@ -786,6 +789,8 @@ class UtSpringEntityManagerModel : UtCustomModel(
     classId = SpringModelUtils.entityManagerClassIds.first(),
     modelName = "entityManager"
 ) {
+    override val dependencies: Collection<UtModel> get() = emptySet()
+
     // NOTE that overriding equals is required just because without it
     // we will lose equality for objects after deserialization
     override fun equals(other: Any?): Boolean = other is UtSpringEntityManagerModel
@@ -814,7 +819,9 @@ data class UtSpringMockMvcResultActionsModel(
     classId = SpringModelUtils.resultActionsClassId,
     id = id,
     modelName = "mockMvcResultActions@$id"
-)
+) {
+    override val dependencies: Collection<UtModel> get() = emptySet()
+}
 
 /**
  * Model for a step to obtain [UtAssembleModel].

utbot-framework/src/main/kotlin/org/utbot/framework/util/UtConcreteExecutionResultUtils.kt

@@ -21,12 +21,9 @@ private fun UtConcreteExecutionResult.updateWithAssembleModels(
     val resolvedResult =
         (result as? UtExecutionSuccess)?.model?.let { UtExecutionSuccess(toAssemble(it)) } ?: result
 
-    return UtConcreteExecutionResult(
-        stateBefore,
-        resolvedStateAfter,
-        resolvedResult,
-        coverage,
-        newInstrumentation
+    return copy(
+        stateAfter = resolvedStateAfter,
+        result = resolvedResult,
     )
 }
 

utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/execution/SimpleUtExecutionInstrumentation.kt

@@ -2,6 +2,7 @@ package org.utbot.instrumentation.instrumentation.execution
 
 import org.utbot.framework.plugin.api.EnvironmentModels
 import org.utbot.framework.plugin.api.FieldId
+import org.utbot.framework.plugin.api.MethodId
 import org.utbot.framework.plugin.api.UtModel
 import org.utbot.framework.plugin.api.util.executable
 import org.utbot.framework.plugin.api.util.signature
@@ -11,6 +12,7 @@ import org.utbot.instrumentation.instrumentation.ArgumentList
 import org.utbot.instrumentation.instrumentation.InvokeInstrumentation
 import org.utbot.instrumentation.instrumentation.et.TraceHandler
 import org.utbot.instrumentation.instrumentation.execution.constructors.ConstructOnlyUserClassesOrCachedObjectsStrategy
+import org.utbot.instrumentation.instrumentation.execution.constructors.StateBeforeAwareIdGenerator
 import org.utbot.instrumentation.instrumentation.execution.constructors.UtModelConstructor
 import org.utbot.instrumentation.instrumentation.execution.context.InstrumentationContext
 import org.utbot.instrumentation.instrumentation.execution.context.SimpleInstrumentationContext
@@ -53,12 +55,17 @@ class SimpleUtExecutionInstrumentation(
         }
         val (stateBefore, instrumentations, timeout) = parameters // smart cast to UtConcreteExecutionData
 
+        lateinit var detectedMockingCandidates: Set<MethodId>
+
         return PhasesController(
             instrumentationContext,
             traceHandler,
             delegateInstrumentation,
-            timeout
+            timeout,
+            idGenerator = StateBeforeAwareIdGenerator.fromUtConcreteExecutionData(parameters)
         ).computeConcreteExecutionResult {
+            detectedMockingCandidates = valueConstructionPhase.detectedMockingCandidates
+
             phasesWrapper {
                 try {
                     // some preparation actions for concrete execution
@@ -126,7 +133,10 @@ class SimpleUtExecutionInstrumentation(
                     applyPostprocessing()
                 }
             }
-        }.toCompleteUtConcreteExecutionResult(stateBefore = stateBefore)
+        }.toCompleteUtConcreteExecutionResult(
+            stateBefore = stateBefore,
+            detectedMockingCandidates = detectedMockingCandidates,
+        )
     }
 
     override fun getStaticField(fieldId: FieldId): Result<UtModel> =

utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/execution/UtExecutionInstrumentation.kt

@@ -24,21 +24,25 @@ data class UtConcreteExecutionData(
 
 /**
  * [UtConcreteExecutionResult] that has not yet been populated with extra data, e.g.:
- *  - updated `stateBefore: EnvironmentModels`
- *  - `detectedMockingCandidates: Set<ExecutableId>` (not yet implemented, see #2321)
+ *  - updated [UtConcreteExecutionResult.stateBefore]
+ *  - [UtConcreteExecutionResult.detectedMockingCandidates]
  */
 data class PreliminaryUtConcreteExecutionResult(
     val stateAfter: EnvironmentModels,
     val result: UtExecutionResult,
     val coverage: Coverage,
     val newInstrumentation: List<UtInstrumentation>? = null,
 ) {
-    fun toCompleteUtConcreteExecutionResult(stateBefore: EnvironmentModels) = UtConcreteExecutionResult(
-        stateBefore,
-        stateAfter,
-        result,
-        coverage,
-        newInstrumentation,
+    fun toCompleteUtConcreteExecutionResult(
+        stateBefore: EnvironmentModels,
+        detectedMockingCandidates: Set<MethodId>
+    ) = UtConcreteExecutionResult(
+        stateBefore = stateBefore,
+        stateAfter = stateAfter,
+        result = result,
+        coverage = coverage,
+        newInstrumentation = newInstrumentation,
+        detectedMockingCandidates = detectedMockingCandidates,
     )
 }
 
@@ -48,6 +52,7 @@ data class UtConcreteExecutionResult(
     val result: UtExecutionResult,
     val coverage: Coverage,
     val newInstrumentation: List<UtInstrumentation>? = null,
+    val detectedMockingCandidates: Set<MethodId>,
 ) {
     override fun toString(): String = buildString {
         appendLine("UtConcreteExecutionResult(")

utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/execution/constructors/InstrumentationContextAwareValueConstructor.kt

@@ -30,25 +30,32 @@ import org.utbot.framework.plugin.api.UtReferenceModel
 import org.utbot.framework.plugin.api.UtStatementCallModel
 import org.utbot.framework.plugin.api.UtStaticMethodInstrumentation
 import org.utbot.framework.plugin.api.UtVoidModel
+import org.utbot.framework.plugin.api.isNull
 import org.utbot.framework.plugin.api.util.anyInstance
+import org.utbot.framework.plugin.api.util.classClassId
 import org.utbot.framework.plugin.api.util.constructor
 import org.utbot.framework.plugin.api.util.constructor.CapturedArgument
 import org.utbot.framework.plugin.api.util.constructor.constructLambda
 import org.utbot.framework.plugin.api.util.constructor.constructStaticLambda
+import org.utbot.framework.plugin.api.util.defaultValueModel
 import org.utbot.framework.plugin.api.util.executableId
+import org.utbot.framework.plugin.api.util.id
 import org.utbot.framework.plugin.api.util.isStatic
 import org.utbot.framework.plugin.api.util.jClass
 import org.utbot.framework.plugin.api.util.jField
 import org.utbot.framework.plugin.api.util.method
+import org.utbot.framework.plugin.api.util.stringClassId
 import org.utbot.framework.plugin.api.util.utContext
-import org.utbot.instrumentation.instrumentation.execution.mock.InstanceMockController
 import org.utbot.instrumentation.instrumentation.execution.context.InstrumentationContext
+import org.utbot.instrumentation.instrumentation.execution.mock.InstanceMockController
 import org.utbot.instrumentation.instrumentation.execution.mock.MethodMockController
 import org.utbot.instrumentation.instrumentation.execution.mock.MockController
 import org.utbot.instrumentation.process.runSandbox
+import java.lang.reflect.Method
 import java.lang.reflect.Modifier
+import java.security.AccessController
+import java.security.PrivilegedAction
 import java.util.*
-import org.utbot.framework.plugin.api.util.id
 import kotlin.reflect.KClass
 
 /**
@@ -63,7 +70,8 @@ import kotlin.reflect.KClass
  */
 // TODO: JIRA:1379 -- Refactor ValueConstructor and InstrumentationContextAwareValueConstructor
 class InstrumentationContextAwareValueConstructor(
-    private val instrumentationContext: InstrumentationContext
+    private val instrumentationContext: InstrumentationContext,
+    private val idGenerator: StateBeforeAwareIdGenerator,
 ) {
     private val classLoader: ClassLoader
         get() = utContext.classLoader
@@ -77,6 +85,8 @@ class InstrumentationContextAwareValueConstructor(
             return objectToModel
         }
 
+    val detectedMockingCandidates: MutableSet<MethodId> = mutableSetOf()
+
     // TODO: JIRA:1379 -- replace UtReferenceModel with Int
     private val constructedObjects = HashMap<UtReferenceModel, Any>()
 
@@ -146,7 +156,7 @@ class InstrumentationContextAwareValueConstructor(
             notMockInstance
         } else {
             val concreteValues = model.mocks.mapValues { mutableListOf<Any?>() }
-            val mockInstance = generateMockitoMock(javaClass, concreteValues)
+            val mockInstance = generateMockitoMock(javaClass, concreteValues, model)
 
             constructedObjects[model] = mockInstance
 
@@ -157,6 +167,12 @@ class InstrumentationContextAwareValueConstructor(
                 valuesList.addAll(constructedValues)
             }
 
+            if (model.canHaveRedundantOrMissingMocks) {
+                // we clear `mocks` to avoid redundant mocks,
+                // actually useful mocks should be later added back as they are used
+                model.mocks.clear()
+            }
+
             mockInstance
         }
 
@@ -194,8 +210,87 @@ class InstrumentationContextAwareValueConstructor(
         }
     }
 
-    private fun generateMockitoMock(clazz: Class<*>, concreteValues: Map<ExecutableId, List<Any?>>): Any {
-        val answer = generateMockitoAnswer(concreteValues)
+    private fun generateMockitoAnswerHandlingRedundantAndMissingMocks(
+        concreteValues: Map<ExecutableId, List<Any?>>,
+        mockModel: UtCompositeModel
+    ): Answer<*> {
+        class MockedExecutable(
+            val executableId: ExecutableId,
+            val answerValues: List<Any?>,
+            val answerModels: List<UtModel>,
+        ) {
+            private var pointer: Int = 0
+
+            fun nextAnswer(): Any? {
+                val answerValue = answerValues[pointer]
+                val answerModel = answerModels[pointer]
+                pointer = (pointer + 1) % answerValues.size
+                (mockModel.mocks.getOrPut(executableId) { mutableListOf() } as MutableList).add(answerModel)
+                return answerValue
+            }
+        }
+
+        val mockedExecutables = concreteValues.mapValues { (executableId, values) ->
+            MockedExecutable(
+                executableId = executableId,
+                answerValues = values,
+                answerModels = mockModel.mocks.getValue(executableId),
+            )
+        }.toMutableMap()
+
+        return Answer { invocation ->
+            with(invocation.method) {
+                mockedExecutables.getOrPut(executableId) {
+                    detectedMockingCandidates.add(executableId)
+                    val answerModel = generateNewAnswerModel()
+
+                    MockedExecutable(
+                        executableId = executableId,
+                        answerValues = listOf(
+                            // `construct()` heavily uses reflection and Mockito,
+                            // so it can't run in sandbox and we need to elevate permissions
+                            AccessController.doPrivileged(PrivilegedAction { construct(answerModel) })
+                                .value.takeUnless { it == Unit }
+                        ),
+                        answerModels = listOf(answerModel)
+                    )
+                }.nextAnswer()
+            }
+        }
+    }
+
+    private fun Method.generateNewAnswerModel() =
+        executableId.returnType.defaultValueModel().takeUnless { it.isNull() } ?: when (executableId.returnType) {
+            // mockito can't mock `String` and `Class`
+            stringClassId -> UtNullModel(stringClassId)
+            classClassId -> UtClassRefModel(
+                id = idGenerator.createId(),
+                classId = classClassId,
+                value = classClassId,
+            )
+
+            else -> UtCompositeModel(
+                id = idGenerator.createId(),
+                // TODO mockito can't mock sealed interfaces,
+                //  we have to mock their implementations or use null
+                classId = executableId.returnType,
+                isMock = true,
+                canHaveRedundantOrMissingMocks = true,
+            )
+        }
+
+    private fun generateMockitoMock(
+        clazz: Class<*>,
+        concreteValues: Map<ExecutableId, List<Any?>>,
+        mockModel: UtCompositeModel
+    ): Any {
+        val answer =
+            if (mockModel.canHaveRedundantOrMissingMocks) {
+                generateMockitoAnswerHandlingRedundantAndMissingMocks(concreteValues, mockModel)
+            } else {
+                generateMockitoAnswer(concreteValues)
+            }
+
         return Mockito.mock(clazz, answer)
     }
 

utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/execution/constructors/UtModelConstructor.kt

@@ -33,18 +33,13 @@ interface UtModelConstructorInterface {
  */
 class UtModelConstructor(
     private val objectToModelCache: IdentityHashMap<Any, UtModel>,
+    private val idGenerator: StateBeforeAwareIdGenerator,
     private val utModelWithCompositeOriginConstructorFinder: (ClassId) -> UtModelWithCompositeOriginConstructor?,
     private val compositeModelStrategy: UtCompositeModelStrategy = AlwaysConstructStrategy,
     private val maxDepth: Long = DEFAULT_MAX_DEPTH
 ) : UtModelConstructorInterface {
     private val constructedObjects = IdentityHashMap<Any, UtModel>()
 
-    private var unusedId = 0
-    private val usedIds = objectToModelCache.values
-        .filterIsInstance<UtReferenceModel>()
-        .mapNotNull { it.id }
-        .toMutableSet()
-
     companion object {
         private const val DEFAULT_MAX_DEPTH = 7L
 
@@ -56,16 +51,16 @@ class UtModelConstructor(
             val strategy = ConstructOnlyUserClassesOrCachedObjectsStrategy(
                 pathsToUserClasses, cache
             )
-            return UtModelConstructor(cache, utModelWithCompositeOriginConstructorFinder, strategy)
+            return UtModelConstructor(
+                objectToModelCache = cache,
+                idGenerator = StateBeforeAwareIdGenerator(preExistingModels = emptySet()),
+                utModelWithCompositeOriginConstructorFinder = utModelWithCompositeOriginConstructorFinder,
+                compositeModelStrategy = strategy
+            )
         }
     }
 
-    private fun computeUnusedIdAndUpdate(): Int {
-        while (unusedId in usedIds) {
-            unusedId++
-        }
-        return unusedId.also { usedIds += it }
-    }
+    private fun computeUnusedIdAndUpdate(): Int = idGenerator.createId()
 
     private fun handleId(value: Any): Int {
         return objectToModelCache[value]?.let { (it as? UtReferenceModel)?.id } ?: computeUnusedIdAndUpdate()

utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/execution/phases/ModelConstructionPhase.kt

@@ -8,6 +8,7 @@ import org.utbot.framework.plugin.api.visible.UtStreamConsumingException
 import org.utbot.instrumentation.instrumentation.et.ExplicitThrowInstruction
 import org.utbot.instrumentation.instrumentation.et.TraceHandler
 import org.utbot.instrumentation.instrumentation.execution.PreliminaryUtConcreteExecutionResult
+import org.utbot.instrumentation.instrumentation.execution.constructors.StateBeforeAwareIdGenerator
 import org.utbot.instrumentation.instrumentation.execution.constructors.UtCompositeModelStrategy
 import org.utbot.instrumentation.instrumentation.execution.constructors.UtModelWithCompositeOriginConstructor
 import org.utbot.instrumentation.instrumentation.execution.constructors.UtModelConstructor
@@ -20,6 +21,7 @@ import java.util.*
 class ModelConstructionPhase(
     private val traceHandler: TraceHandler,
     private val utModelWithCompositeOriginConstructorFinder: (ClassId) -> UtModelWithCompositeOriginConstructor?,
+    private val idGenerator: StateBeforeAwareIdGenerator,
 ) : ExecutionPhase {
 
     override fun wrapError(e: Throwable): ExecutionPhaseException {
@@ -52,6 +54,7 @@ class ModelConstructionPhase(
                 objectToModelCache = cache,
                 utModelWithCompositeOriginConstructorFinder = utModelWithCompositeOriginConstructorFinder,
                 compositeModelStrategy = strategy,
+                idGenerator = idGenerator,
             )
         }
     }