Avoid harcoding the names of tables and repositories

Author: EgorkaKulikov

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/domain/builtin/SpringBuiltins.kt

@@ -5,4 +5,4 @@ import org.utbot.framework.plugin.api.BuiltinClassId
 internal val autowiredClassId = BuiltinClassId(
     canonicalName = "org.springframework.beans.factory.annotation.Autowired",
     simpleName = "Autowired",
-)
+)

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

@@ -6,9 +6,13 @@ import org.utbot.instrumentation.instrumentation.ArgumentList
 import org.utbot.instrumentation.instrumentation.Instrumentation
 import org.utbot.instrumentation.instrumentation.execution.mock.SpringInstrumentationContext
 import org.utbot.instrumentation.process.HandlerClassesLoader
+import org.utbot.instrumentation.util.AccessibleTypesAnalyzer
 import org.utbot.spring.api.context.ContextWrapper
+import org.utbot.spring.api.context.NamedBean
 import org.utbot.spring.api.repositoryWrapper.RepositoryInteraction
+import java.lang.reflect.ParameterizedType
 import java.security.ProtectionDomain
+import java.util.IdentityHashMap
 
 /**
  * UtExecutionInstrumentation wrapper that is aware of Spring config and initialises Spring context
@@ -20,6 +24,9 @@ class SpringUtExecutionInstrumentation(
     private lateinit var instrumentationContext: SpringInstrumentationContext
     private val springContext: ContextWrapper get() = instrumentationContext.springContext
 
+    private val typesAnalyzer = AccessibleTypesAnalyzer()
+    private val relatedBeansCache = IdentityHashMap<Class<*>, Set<NamedBean>>()
+
     companion object {
         private val logger = getLogger<SpringUtExecutionInstrumentation>()
         private const val SPRING_COMMONS_JAR_FILENAME = "utbot-spring-commons-shadow.jar"
@@ -43,22 +50,39 @@ class SpringUtExecutionInstrumentation(
         parameters: Any?
     ): UtConcreteExecutionResult {
         RepositoryInteraction.recordedInteractions.clear()
-        // TODO properly detect which beans need to be reset, right now "orderRepository" and "orderService" are hardcoded
-        val beanNamesToReset = listOf("orderRepository", "orderService")
 
-        beanNamesToReset.forEach { beanNameToReset ->
-            val beanDefToReset = springContext.getBeanDefinition(beanNameToReset)
-            springContext.removeBeanDefinition(beanNameToReset)
-            springContext.registerBeanDefinition(beanNameToReset, beanDefToReset)
+        val beansToReset: Set<NamedBean> = relatedBeansCache.getOrPut(clazz) {
+            val relatedBeans = springContext.getBeansOfType(clazz)
+            val accessibleTypes = relatedBeans.flatMap { typesAnalyzer.collectAccessibleTypes(it.bean) }
+
+            val beansToReset = accessibleTypes
+                .flatMap { springContext.getBeansOfType(it) }
+                .toSet()
+
+            beansToReset
         }
 
+        beansToReset.forEach { springContext.resetBean(it.beanName) }
         val jdbcTemplate = getBean("jdbcTemplate")
-        // TODO properly detect which repositories need to be cleared, right now "orders" is hardcoded
-        val sql = "TRUNCATE TABLE orders"
+
+        val entityNames = beansToReset
+            .filter {
+                val beanType =  it.bean::class.java
+                val implementedInterfaces = beanType.genericInterfaces.map { i -> i.javaClass.canonicalName }
+                "org.springframework.data.repository.Repository" in  implementedInterfaces
+            }
+            .map { (it.bean::class.java.genericInterfaces[0] as ParameterizedType).actualTypeArguments[0] }
+            .map { it.typeName }
+
+        //TODO: table name may be different from entity name (set with annotation on entity)
+        //TODO: plural form is not always obtained with "+s", it can be even Person->People
+        val tableNames = entityNames.map { it + "s" }
+
+        val sql = "TRUNCATE TABLE $tableNames"
         jdbcTemplate::class.java
             .getMethod("execute", sql::class.java)
             .invoke(jdbcTemplate, sql)
-        val sql2 = "ALTER TABLE orders ALTER COLUMN id RESTART WITH 1"
+        val sql2 = "ALTER TABLE $tableNames ALTER COLUMN id RESTART WITH 1"
         jdbcTemplate::class.java
             .getMethod("execute", sql::class.java)
             .invoke(jdbcTemplate, sql2)

utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/util/AccessibleTypesAnalyzer.kt

@@ -0,0 +1,36 @@
+package org.utbot.instrumentation.util
+
+import java.util.IdentityHashMap
+
+class AccessibleTypesAnalyzer {
+    fun collectAccessibleTypes(bean: Any): Set<Class<*>> {
+        val accessibleObjects = collectFromObject(bean, analyzedObjects = IdentityHashMap())
+
+        return accessibleObjects
+            .map { it::class.java }
+            .toSet()
+    }
+
+    private fun collectFromObject(obj: Any, analyzedObjects: IdentityHashMap<Any, Unit>): Set<Any> {
+        if (analyzedObjects.contains(obj)) {
+            return emptySet()
+        }
+
+        analyzedObjects[obj] = Unit
+
+        val clazz = obj::class.java
+        val objects = mutableSetOf<Any>()
+
+        var current: Class<*> = clazz
+        while (current.superclass != null) {
+            objects.addAll(current.fields.mapNotNull {
+                it.get(obj)
+            }.flatMap {
+                listOf(it) + collectFromObject(it, analyzedObjects)
+            })
+            current = current.superclass
+        }
+
+        return objects
+    }
+}

utbot-spring-commons-api/src/main/kotlin/org/utbot/spring/api/context/ContextWrapper.kt

@@ -3,11 +3,14 @@ package org.utbot.spring.api.context
 interface ContextWrapper {
     val context: Any
 
-    fun getBeanDefinition(beanName: String): Any
+    fun getBean(beanName: String): Any
 
-    fun removeBeanDefinition(beanName: String)
+    fun getBeansOfType(clazz: Class<*>): Set<NamedBean>
 
-    fun registerBeanDefinition(beanName: String, beanDefinition: Any)
+    fun resetBean(beanName: String): Any
+}
 
-    fun getBean(beanName: String): Any
-}
+data class NamedBean(
+    val beanName: String,
+    val bean: Any,
+)

utbot-spring-commons-api/src/main/kotlin/org/utbot/spring/api/repositoryWrapper/RepositoryInteraction.kt

@@ -11,4 +11,4 @@ data class RepositoryInteraction(
     companion object {
         val recordedInteractions = mutableListOf<RepositoryInteraction>()
     }
-}
+}

utbot-spring-commons/src/main/kotlin/org/utbot/spring/context/SpringContextWrapper.kt

@@ -1,23 +1,21 @@
 package org.utbot.spring.context
 
-import org.springframework.beans.factory.config.BeanDefinition
 import org.springframework.beans.factory.support.BeanDefinitionRegistry
 import org.springframework.context.ConfigurableApplicationContext
 import org.utbot.spring.api.context.ContextWrapper
+import org.utbot.spring.api.context.NamedBean
 
 class SpringContextWrapper(override val context: ConfigurableApplicationContext) : ContextWrapper {
-    override fun getBeanDefinition(beanName: String): BeanDefinition = context.beanFactory.getBeanDefinition(beanName)
+    override fun getBean(beanName: String): Any = context.getBean(beanName)
 
-    override fun removeBeanDefinition(beanName: String) {
-        val beanDefinitionRegistry = context.beanFactory as BeanDefinitionRegistry
-        beanDefinitionRegistry.removeBeanDefinition(beanName)
-    }
+    override fun getBeansOfType(clazz: Class<*>): Set<NamedBean> =
+        context.beanFactory.getBeansOfType(clazz).map { NamedBean(it.key, it.value) }.toSet()
 
-    override fun registerBeanDefinition(beanName: String, beanDefinition: Any) {
+    override fun resetBean(beanName: String) {
         val beanDefinitionRegistry = context.beanFactory as BeanDefinitionRegistry
-        beanDefinition as BeanDefinition
+
+        val beanDefinition = context.beanFactory.getBeanDefinition(beanName)
+        beanDefinitionRegistry.removeBeanDefinition(beanName)
         beanDefinitionRegistry.registerBeanDefinition(beanName, beanDefinition)
     }
-
-    override fun getBean(beanName: String): Any = context.getBean(beanName)
 }