Automatically detect the names of tables and repositories

utbot-framework/src/main/kotlin/org/utbot/engine/UtBotSymbolicEngine.kt

@@ -46,6 +46,7 @@ import org.utbot.fuzzing.spring.SavedEntityProvider
 import org.utbot.fuzzing.spring.SpringBeanValueProvider
 import org.utbot.fuzzing.utils.Trie
 import org.utbot.instrumentation.ConcreteExecutor
+import org.utbot.instrumentation.getRelevantSpringRepositories
 import org.utbot.instrumentation.instrumentation.Instrumentation
 import org.utbot.instrumentation.instrumentation.execution.UtConcreteExecutionData
 import org.utbot.instrumentation.instrumentation.execution.UtConcreteExecutionResult
@@ -388,18 +389,12 @@ class UtBotSymbolicEngine(
             .letIf(applicationContext is SpringApplicationContext
                         && applicationContext.typeReplacementApproach is TypeReplacementApproach.ReplaceIfPossible
             ) { provider ->
-                // TODO #2274 properly detect relevant repositories (right now orderRepository is hardcoded)
-                val relevantRepositories = listOf(
-                    SpringRepositoryId(
-                        repositoryBeanName = "orderRepository",
-                        repositoryClassId = ClassId("com.rest.order.repositories.OrderRepository"),
-                        entityClassId = ClassId("com.rest.order.models.Order")
-                    )
-                )
-                val generatedValueFieldIds = listOf(
-                    FieldId(ClassId("com.rest.order.models.Order"), "id")
-                )
-                logger.info { "Relevant repositories: $relevantRepositories" }
+                val relevantRepositories = concreteExecutor.getRelevantSpringRepositories(methodUnderTest.classId)
+                val generatedValueFieldIds = relevantRepositories.map {
+                    repository -> FieldId(repository.entityClassId, "id")
+                }
+
+                logger.info { "Detected relevant repositories for class ${methodUnderTest.classId}: $relevantRepositories" }
                 // spring should try to generate bean values, but if it fails, then object value provider is used for it
                 val springBeanValueProvider = SpringBeanValueProvider(
                     defaultIdGenerator,

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-framework/src/main/kotlin/org/utbot/framework/plugin/api/TestCaseGenerator.kt

@@ -73,7 +73,11 @@ open class TestCaseGenerator(
                 is TypeReplacementApproach.ReplaceIfPossible ->
                     when (applicationContext.testType) {
                         SpringTestsType.UNIT_TESTS -> UtExecutionInstrumentation
-                        SpringTestsType.INTEGRATION_TESTS -> SpringUtExecutionInstrumentation(UtExecutionInstrumentation, approach.config)
+                        SpringTestsType.INTEGRATION_TESTS -> SpringUtExecutionInstrumentation(
+                                UtExecutionInstrumentation,
+                                approach.config,
+                                applicationContext.beanDefinitions,
+                                )
                     }
                 is TypeReplacementApproach.DoNotReplace -> UtExecutionInstrumentation
             }

utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/ConcreteExecutor.kt

@@ -1,9 +1,5 @@
 package org.utbot.instrumentation
 
-import com.jetbrains.rd.util.ILoggerFactory
-import com.jetbrains.rd.util.Logger
-import com.jetbrains.rd.util.Statics
-import com.jetbrains.rd.util.lifetime.Lifetime
 import com.jetbrains.rd.util.lifetime.LifetimeDefinition
 import com.jetbrains.rd.util.lifetime.isAlive
 import com.jetbrains.rd.util.lifetime.throwIfNotAlive
@@ -15,24 +11,25 @@ import kotlin.reflect.KProperty
 import kotlin.reflect.jvm.javaConstructor
 import kotlin.reflect.jvm.javaGetter
 import kotlin.reflect.jvm.javaMethod
-import kotlin.streams.toList
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
 import mu.KotlinLogging
 import org.utbot.framework.plugin.api.InstrumentedProcessDeathException
 import org.utbot.common.logException
+import org.utbot.framework.plugin.api.ClassId
 import org.utbot.framework.plugin.api.FieldId
+import org.utbot.framework.plugin.api.SpringRepositoryId
 import org.utbot.framework.plugin.api.util.UtContext
 import org.utbot.framework.plugin.api.util.signature
 import org.utbot.instrumentation.instrumentation.Instrumentation
 import org.utbot.instrumentation.process.generated.ComputeStaticFieldParams
+import org.utbot.instrumentation.process.generated.GetSpringRepositoriesParams
 import org.utbot.instrumentation.process.generated.InvokeMethodCommandParams
 import org.utbot.instrumentation.rd.InstrumentedProcess
 import org.utbot.instrumentation.util.InstrumentedProcessError
 import org.utbot.rd.generated.synchronizationModel
-import org.utbot.rd.loggers.UtRdKLoggerFactory
 import org.utbot.rd.loggers.overrideDefaultRdLoggerFactoryWithKLogger
 
 private val logger = KotlinLogging.logger {}
@@ -285,6 +282,16 @@ fun ConcreteExecutor<*,*>.warmup() = runBlocking {
     }
 }
 
+fun ConcreteExecutor<*, *>.getRelevantSpringRepositories(classId: ClassId): Set<SpringRepositoryId> = runBlocking {
+    withProcess {
+        val classId = kryoHelper.writeObject(classId)
+        val params = GetSpringRepositoriesParams(classId)
+        val result = instrumentedProcessModel.getRelevantSpringRepositories.startSuspending(lifetime, params)
+
+        kryoHelper.readObject(result.springRepositoryIds)
+    }
+}
+
 /**
  * Extension function for the [ConcreteExecutor], which allows to collect static field value of [fieldId].
  */

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

@@ -2,22 +2,32 @@ package org.utbot.instrumentation.instrumentation.execution
 
 import org.utbot.common.JarUtils
 import com.jetbrains.rd.util.getLogger
+import com.jetbrains.rd.util.info
+import org.utbot.framework.plugin.api.BeanDefinitionData
+import org.utbot.framework.plugin.api.ClassId
+import org.utbot.framework.plugin.api.SpringRepositoryId
+import org.utbot.framework.plugin.api.util.jClass
 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.spring.api.context.ContextWrapper
 import org.utbot.spring.api.repositoryWrapper.RepositoryInteraction
 import java.security.ProtectionDomain
+import java.util.IdentityHashMap
 
 /**
  * UtExecutionInstrumentation wrapper that is aware of Spring config and initialises Spring context
  */
 class SpringUtExecutionInstrumentation(
     private val delegateInstrumentation: UtExecutionInstrumentation,
-    private val springConfig: String
+    private val springConfig: String,
+    private val beanDefinitions: List<BeanDefinitionData>,
 ) : Instrumentation<UtConcreteExecutionResult> by delegateInstrumentation {
     private lateinit var instrumentationContext: SpringInstrumentationContext
+
+    private val relatedBeansCache = mutableMapOf<Class<*>, Set<String>>()
+
     private val springContext: ContextWrapper get() = instrumentationContext.springContext
 
     companion object {
@@ -26,11 +36,16 @@ class SpringUtExecutionInstrumentation(
     }
 
     override fun init(pathsToUserClasses: Set<String>) {
-        HandlerClassesLoader.addUrls(listOf(JarUtils.extractJarFileFromResources(
-            jarFileName = SPRING_COMMONS_JAR_FILENAME,
-            jarResourcePath = "lib/$SPRING_COMMONS_JAR_FILENAME",
-            targetDirectoryName = "spring-commons"
-        ).path))
+        HandlerClassesLoader.addUrls(
+            listOf(
+                JarUtils.extractJarFileFromResources(
+                    jarFileName = SPRING_COMMONS_JAR_FILENAME,
+                    jarResourcePath = "lib/$SPRING_COMMONS_JAR_FILENAME",
+                    targetDirectoryName = "spring-commons"
+                ).path
+            )
+        )
+
         instrumentationContext = SpringInstrumentationContext(springConfig)
         delegateInstrumentation.instrumentationContext = instrumentationContext
         delegateInstrumentation.init(pathsToUserClasses)
@@ -43,39 +58,48 @@ 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 beanNamesToReset: Set<String> = getRelevantBeanNames(clazz)
+        val repositoryDefinitions = springContext.resolveRepositories(beanNamesToReset)
 
+        beanNamesToReset.forEach { beanName -> springContext.resetBean(beanName) }
         val jdbcTemplate = getBean("jdbcTemplate")
-        // TODO properly detect which repositories need to be cleared, right now "orders" is hardcoded
-        val sql = "TRUNCATE TABLE orders"
-        jdbcTemplate::class.java
-            .getMethod("execute", sql::class.java)
-            .invoke(jdbcTemplate, sql)
-        val sql2 = "ALTER TABLE orders ALTER COLUMN id RESTART WITH 1"
-        jdbcTemplate::class.java
-            .getMethod("execute", sql::class.java)
-            .invoke(jdbcTemplate, sql2)
+
+        for (repositoryDefinition in repositoryDefinitions) {
+            val truncateTableCommand = "TRUNCATE TABLE ${repositoryDefinition.tableName}"
+            jdbcTemplate::class.java
+                .getMethod("execute", truncateTableCommand::class.java)
+                .invoke(jdbcTemplate, truncateTableCommand)
+
+            val restartIdCommand = "ALTER TABLE ${repositoryDefinition.tableName} ALTER COLUMN id RESTART WITH 1"
+            jdbcTemplate::class.java
+                .getMethod("execute", restartIdCommand::class.java)
+                .invoke(jdbcTemplate, restartIdCommand)
+        }
 
         return delegateInstrumentation.invoke(clazz, methodSignature, arguments, parameters)
     }
 
+    private fun getRelevantBeanNames(clazz: Class<*>): Set<String> = relatedBeansCache.getOrPut(clazz) {
+        beanDefinitions
+            .filter { it.beanTypeFqn == clazz.name }
+            .flatMap { springContext.getDependenciesForBean(it.beanName) }
+            .toSet()
+            .also { logger.info { "Detected relevant beans for class ${clazz.name}: $it" } }
+    }
+
     fun getBean(beanName: String): Any = springContext.getBean(beanName)
 
-    fun saveToRepository(repository: Any, entity: Any) {
-        // ignore repository interactions done during repository fill up
-        val savedRecordedRepositoryResponses = RepositoryInteraction.recordedInteractions.toList()
-        repository::class.java
-            .getMethod("save", Any::class.java)
-            .invoke(repository, entity)
-        RepositoryInteraction.recordedInteractions.clear()
-        RepositoryInteraction.recordedInteractions.addAll(savedRecordedRepositoryResponses)
+    fun getRepositoryDescriptions(classId: ClassId): Set<SpringRepositoryId> {
+        val relevantBeanNames = getRelevantBeanNames(classId.jClass)
+        val repositoryDescriptions = springContext.resolveRepositories(relevantBeanNames.toSet())
+        return repositoryDescriptions.map { repositoryDescription ->
+            SpringRepositoryId(
+                repositoryDescription.beanName,
+                ClassId(repositoryDescription.repositoryName),
+                ClassId(repositoryDescription.entityName),
+            )
+        }.toSet()
     }
 
     override fun transform(
@@ -85,14 +109,15 @@ class SpringUtExecutionInstrumentation(
         protectionDomain: ProtectionDomain,
         classfileBuffer: ByteArray
     ): ByteArray? =
-        // TODO automatically detect which libraries we don't want to transform (by total transformation time)
-        // transforming Spring takes too long
+        // TODO: automatically detect which libraries we don't want to transform (by total transformation time)
         if (listOf(
                 "org/springframework",
                 "com/fasterxml",
                 "org/hibernate",
                 "org/apache",
-                "org/h2"
+                "org/h2",
+                "javax/",
+                "ch/qos",
             ).any { className.startsWith(it) }
         ) {
             null

utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/process/InstrumentedProcessMain.kt

@@ -15,6 +15,7 @@ import org.utbot.instrumentation.instrumentation.execution.SpringUtExecutionInst
 import org.utbot.instrumentation.instrumentation.execution.constructors.UtModelConstructor
 import org.utbot.instrumentation.process.generated.CollectCoverageResult
 import org.utbot.instrumentation.process.generated.GetSpringBeanResult
+import org.utbot.instrumentation.process.generated.GetSpringRepositoriesResult
 import org.utbot.instrumentation.process.generated.InstrumentedProcessModel
 import org.utbot.instrumentation.process.generated.InvokeMethodCommandResult
 import org.utbot.instrumentation.process.generated.instrumentedProcessModel
@@ -167,4 +168,9 @@ private fun InstrumentedProcessModel.setup(kryoHelper: KryoHelper, watchdog: Idl
         )
         GetSpringBeanResult(kryoHelper.writeObject(model))
     }
+    watchdog.measureTimeForActiveCall(getRelevantSpringRepositories, "Getting Spring repositories") { params ->
+        val classId: ClassId = kryoHelper.readObject(params.classId)
+        val repositoryDescriptions = (instrumentation as SpringUtExecutionInstrumentation).getRepositoryDescriptions(classId)
+        GetSpringRepositoriesResult(kryoHelper.writeObject(repositoryDescriptions))
+    }
 }