@@ -2,22 +2,31 @@ package org.utbot.instrumentation.instrumentation.execution
22
33import org.utbot.common.JarUtils
44import com.jetbrains.rd.util.getLogger
5+ import org.utbot.framework.plugin.api.BeanDefinitionData
6+ import org.utbot.framework.plugin.api.ClassId
7+ import org.utbot.framework.plugin.api.SpringRepositoryId
8+ import org.utbot.framework.plugin.api.util.jClass
59import org.utbot.instrumentation.instrumentation.ArgumentList
610import org.utbot.instrumentation.instrumentation.Instrumentation
711import org.utbot.instrumentation.instrumentation.execution.mock.SpringInstrumentationContext
812import org.utbot.instrumentation.process.HandlerClassesLoader
913import org.utbot.spring.api.context.ContextWrapper
1014import org.utbot.spring.api.repositoryWrapper.RepositoryInteraction
1115import java.security.ProtectionDomain
16+ import java.util.IdentityHashMap
1217
1318/* *
1419 * UtExecutionInstrumentation wrapper that is aware of Spring config and initialises Spring context
1520 */
1621class SpringUtExecutionInstrumentation (
1722 private val delegateInstrumentation : UtExecutionInstrumentation ,
18- private val springConfig : String
23+ private val springConfig : String ,
24+ private val beanDefinitions : List <BeanDefinitionData >,
1925) : Instrumentation<UtConcreteExecutionResult> by delegateInstrumentation {
2026 private lateinit var instrumentationContext: SpringInstrumentationContext
27+
28+ private lateinit var relatedBeansCache: IdentityHashMap <Class <* >, Set <String >>
29+
2130 private val springContext: ContextWrapper get() = instrumentationContext.springContext
2231
2332 companion object {
@@ -26,14 +35,21 @@ class SpringUtExecutionInstrumentation(
2635 }
2736
2837 override fun init (pathsToUserClasses : Set <String >) {
29- HandlerClassesLoader .addUrls(listOf (JarUtils .extractJarFileFromResources(
30- jarFileName = SPRING_COMMONS_JAR_FILENAME ,
31- jarResourcePath = " lib/$SPRING_COMMONS_JAR_FILENAME " ,
32- targetDirectoryName = " spring-commons"
33- ).path))
38+ HandlerClassesLoader .addUrls(
39+ listOf (
40+ JarUtils .extractJarFileFromResources(
41+ jarFileName = SPRING_COMMONS_JAR_FILENAME ,
42+ jarResourcePath = " lib/$SPRING_COMMONS_JAR_FILENAME " ,
43+ targetDirectoryName = " spring-commons"
44+ ).path
45+ )
46+ )
47+
3448 instrumentationContext = SpringInstrumentationContext (springConfig)
3549 delegateInstrumentation.instrumentationContext = instrumentationContext
3650 delegateInstrumentation.init (pathsToUserClasses)
51+
52+ relatedBeansCache = IdentityHashMap <Class <* >, Set <String >>()
3753 }
3854
3955 override fun invoke (
@@ -43,39 +59,46 @@ class SpringUtExecutionInstrumentation(
4359 parameters : Any?
4460 ): UtConcreteExecutionResult {
4561 RepositoryInteraction .recordedInteractions.clear()
46- // TODO properly detect which beans need to be reset, right now "orderRepository" and "orderService" are hardcoded
47- val beanNamesToReset = listOf (" orderRepository" , " orderService" )
4862
49- beanNamesToReset.forEach { beanNameToReset ->
50- val beanDefToReset = springContext.getBeanDefinition(beanNameToReset)
51- springContext.removeBeanDefinition(beanNameToReset)
52- springContext.registerBeanDefinition(beanNameToReset, beanDefToReset)
53- }
63+ val beanNamesToReset: Set <String > = relatedBeansCache.getOrPut(clazz) { getRelevantBeanNames(clazz) }
64+ val repositoryDefinitions = springContext.resolveRepositories(beanNamesToReset)
5465
66+ beanNamesToReset.forEach { beanName -> springContext.resetBean(beanName) }
5567 val jdbcTemplate = getBean(" jdbcTemplate" )
56- // TODO properly detect which repositories need to be cleared, right now "orders" is hardcoded
57- val sql = " TRUNCATE TABLE orders"
58- jdbcTemplate::class .java
59- .getMethod(" execute" , sql::class .java)
60- .invoke(jdbcTemplate, sql)
61- val sql2 = " ALTER TABLE orders ALTER COLUMN id RESTART WITH 1"
62- jdbcTemplate::class .java
63- .getMethod(" execute" , sql::class .java)
64- .invoke(jdbcTemplate, sql2)
68+
69+ for (repositoryDefinition in repositoryDefinitions) {
70+ val truncateTableCommand = " TRUNCATE TABLE ${repositoryDefinition.tableName} "
71+ jdbcTemplate::class .java
72+ .getMethod(" execute" , truncateTableCommand::class .java)
73+ .invoke(jdbcTemplate, truncateTableCommand)
74+
75+ val restartIdCommand = " ALTER TABLE ${repositoryDefinition.tableName} ALTER COLUMN id RESTART WITH 1"
76+ jdbcTemplate::class .java
77+ .getMethod(" execute" , restartIdCommand::class .java)
78+ .invoke(jdbcTemplate, restartIdCommand)
79+ }
6580
6681 return delegateInstrumentation.invoke(clazz, methodSignature, arguments, parameters)
6782 }
6883
84+ private fun getRelevantBeanNames (clazz : Class <* >): Set <String > =
85+ beanDefinitions
86+ .filter { it.beanTypeFqn == clazz.name }
87+ .flatMap { springContext.getDependenciesForBean(it.beanName) }
88+ .toSet()
89+
6990 fun getBean (beanName : String ): Any = springContext.getBean(beanName)
7091
71- fun saveToRepository (repository : Any , entity : Any ) {
72- // ignore repository interactions done during repository fill up
73- val savedRecordedRepositoryResponses = RepositoryInteraction .recordedInteractions.toList()
74- repository::class .java
75- .getMethod(" save" , Any ::class .java)
76- .invoke(repository, entity)
77- RepositoryInteraction .recordedInteractions.clear()
78- RepositoryInteraction .recordedInteractions.addAll(savedRecordedRepositoryResponses)
92+ fun getRepositoryDescriptions (classId : ClassId ): Set <SpringRepositoryId > {
93+ val relevantBeanNames = getRelevantBeanNames(classId.jClass)
94+ val repositoryDescriptions = springContext.resolveRepositories(relevantBeanNames.toSet())
95+ return repositoryDescriptions.map { repositoryDescription ->
96+ SpringRepositoryId (
97+ repositoryDescription.beanName,
98+ ClassId (repositoryDescription.repositoryName),
99+ ClassId (repositoryDescription.entityName),
100+ )
101+ }.toSet()
79102 }
80103
81104 override fun transform (
@@ -85,14 +108,15 @@ class SpringUtExecutionInstrumentation(
85108 protectionDomain : ProtectionDomain ,
86109 classfileBuffer : ByteArray
87110 ): ByteArray? =
88- // TODO automatically detect which libraries we don't want to transform (by total transformation time)
89- // transforming Spring takes too long
111+ // TODO: automatically detect which libraries we don't want to transform (by total transformation time)
90112 if (listOf (
91113 " org/springframework" ,
92114 " com/fasterxml" ,
93115 " org/hibernate" ,
94116 " org/apache" ,
95- " org/h2"
117+ " org/h2" ,
118+ " javax/" ,
119+ " ch/qos" ,
96120 ).any { className.startsWith(it) }
97121 ) {
98122 null
0 commit comments