@@ -12,6 +12,9 @@ import org.utbot.framework.plugin.api.UtModel
1212import org.utbot.framework.plugin.api.UtNullModel
1313import org.utbot.framework.plugin.api.UtPrimitiveModel
1414import org.utbot.framework.plugin.api.UtSpringContextModel
15+ import org.utbot.framework.plugin.api.UtStatementModel
16+ import org.utbot.framework.plugin.api.mapper.UtModelDeepMapper
17+ import org.utbot.framework.plugin.api.mapper.mapModels
1518import java.util.Optional
1619
1720object SpringModelUtils {
@@ -108,7 +111,12 @@ object SpringModelUtils {
108111 bypassesSandbox = true // TODO may be we can use some alternative sandbox that has more permissions
109112 )
110113
111- fun createBeanModel (beanName : String , id : Int , classId : ClassId ) = UtAssembleModel (
114+ fun createBeanModel (
115+ beanName : String ,
116+ id : Int ,
117+ classId : ClassId ,
118+ modificationChainProvider : UtAssembleModel .() -> List <UtStatementModel > = { mutableListOf() },
119+ ) = UtAssembleModel (
112120 id = id,
113121 classId = classId,
114122 modelName = " @Autowired $beanName " ,
@@ -117,7 +125,7 @@ object SpringModelUtils {
117125 executable = getBeanMethodId,
118126 params = listOf (UtPrimitiveModel (beanName))
119127 ),
120- modificationsChainProvider = { mutableListOf () }
128+ modificationsChainProvider = modificationChainProvider
121129 )
122130
123131 fun UtModel.isAutowiredFromContext (): Boolean =
@@ -341,8 +349,19 @@ object SpringModelUtils {
341349 returnType = resultMatcherClassId
342350 )
343351
344- fun createMockMvcModel (idGenerator : () -> Int ) =
345- createBeanModel(" mockMvc" , idGenerator(), mockMvcClassId)
352+ fun createMockMvcModel (controller : UtModel ? , idGenerator : () -> Int ) =
353+ createBeanModel(" mockMvc" , idGenerator(), mockMvcClassId, modificationChainProvider = {
354+ // we need to keep controller modifications if there are any, so we add them to mockMvc
355+ (controller as ? UtAssembleModel )?.let { assembledController ->
356+ val controllerModificationRemover = UtModelDeepMapper { model ->
357+ if (model == assembledController) assembledController.copy(modificationsChain = emptyList())
358+ else model
359+ }
360+ // modificationsChain may mention controller, causing controller modifications to evaluate twice:
361+ // once for mockMvc and once for controller itself, to avoid that we remove modifications from controller
362+ assembledController.modificationsChain.map { it.mapModels(controllerModificationRemover) }
363+ } ? : mutableListOf ()
364+ })
346365
347366 fun createRequestBuilderModelOrNull (methodId : MethodId , arguments : List <UtModel >, idGenerator : () -> Int ): UtModel ? {
348367 check(methodId.parameters.size == arguments.size)
0 commit comments