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

@@ -14,10 +14,12 @@ import org.utbot.instrumentation.instrumentation.Instrumentation
 import org.utbot.instrumentation.instrumentation.execution.mock.SpringInstrumentationContext
 import org.utbot.instrumentation.instrumentation.execution.phases.ModelConstructionPhase
 import org.utbot.instrumentation.process.HandlerClassesLoader
+import org.utbot.instrumentation.util.ReflectionUtils
 import org.utbot.spring.api.context.ContextWrapper
 import org.utbot.spring.api.instantiator.InstantiationSettings
 import org.utbot.spring.api.instantiator.ApplicationInstantiatorFacade
 import org.utbot.spring.api.repositoryWrapper.RepositoryInteraction
+import java.lang.reflect.ParameterizedType
 import java.security.ProtectionDomain
 
 /**
@@ -79,22 +81,37 @@ 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 relatedBeanDefinitions = springContext.getRelatedBeanTypes(clazz)
+        val typesFromClass = ReflectionUtils().collectAccessibleTypes(clazz)
+
+        //val beanNamesToReset = listOf("orderRepository", "orderService")
+        val beanNamesToReset = relatedBeanDefinitions
+            .filter { typesFromClass.any { typeFromClass -> it.beanType.isAssignableFrom(typeFromClass) } }
+            .map { it.beanName }
+
+        beanNamesToReset.forEach { beanNameToReset -> springContext.resetBean(beanNameToReset) }
 
         val jdbcTemplate = getBean("jdbcTemplate")
-        // TODO properly detect which repositories need to be cleared, right now "orders" is hardcoded
-        val sql = "TRUNCATE TABLE orders"
+
+        //val tableNames = "orders"
+        val entityNames = relatedBeanDefinitions
+            .filter {
+                val implementedInterfaces = it.beanType.genericInterfaces.map { i -> i.javaClass.canonicalName }
+                "org.springframework.data.repository.Repository" in  implementedInterfaces
+            }
+            .map { (it.beanType.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 $entityNames"
         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 $entityNames 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/ReflectionUtils.kt

@@ -0,0 +1,41 @@
+package org.utbot.instrumentation.util
+
+class ReflectionUtils {
+
+    private val collectedFieldTypes = mutableSetOf<Class<*>>()
+    private val alreadyAnalyzedTypes = mutableSetOf<Class<*>>()
+
+    fun collectAccessibleTypes(clazz: Class<*>): Set<Class<*>> {
+        collectedFieldTypes.clear()
+        alreadyAnalyzedTypes.clear()
+
+        return collectAccessibleFieldTypes(clazz)
+    }
+
+    private fun collectAccessibleFieldTypes(clazz: Class<*>, isTopLevel: Boolean = true): Set<Class<*>> {
+        if (alreadyAnalyzedTypes.contains(clazz)) {
+            return emptySet()
+        }
+
+        collectedFieldTypes += collectFieldTypes(clazz)
+
+        val fieldTypes = clazz.fields.map { it.type }
+        for (fieldType in fieldTypes) {
+            collectedFieldTypes += collectAccessibleFieldTypes(fieldType)
+        }
+
+        return collectedFieldTypes
+    }
+
+    private fun collectFieldTypes(clazz: Class<*>): Set<Class<*>> {
+        val types = mutableSetOf<Class<*>>()
+
+        var current: Class<*> = clazz
+        while (current.superclass != null) {
+            types += current.fields.map { it.type }
+            current = current.superclass
+        }
+
+        return types
+    }
+}

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 getRelatedBeanTypes(clazz: Class<*>): Set<SimpleBeanDefinition>
 
-    fun registerBeanDefinition(beanName: String, beanDefinition: Any)
+    fun resetBean(beanName: String): Any
+}
 
-    fun getBean(beanName: String): Any
-}
+data class SimpleBeanDefinition(
+    val beanName: String,
+    val beanType: Class<*>,
+)

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,30 @@
 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.SimpleBeanDefinition
 
 class SpringContextWrapper(override val context: ConfigurableApplicationContext) : ContextWrapper {
-    override fun getBeanDefinition(beanName: String): BeanDefinition = context.beanFactory.getBeanDefinition(beanName)
-
-    override fun removeBeanDefinition(beanName: String) {
-        val beanDefinitionRegistry = context.beanFactory as BeanDefinitionRegistry
-        beanDefinitionRegistry.removeBeanDefinition(beanName)
+    override fun getBean(beanName: String): Any = context.getBean(beanName)
+    override fun getRelatedBeanTypes(clazz: Class<*>): Set<SimpleBeanDefinition> {
+        val cutName = clazz.canonicalName
+        return context.beanDefinitionNames
+            .map { name -> name to context.getBean(name) }
+            .filter {
+                val bean = it.second
+                val typeName = bean.javaClass.canonicalName
+                typeName != null && cutName.commonPrefixWith(typeName).isNotEmpty()
+            }
+            .map { nameToBean -> SimpleBeanDefinition(nameToBean.first, nameToBean.second.javaClass) }
+            .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)
 }