[split tests] AT-3274 IDE runs as a parameter in tests

GitOrigin-RevId: 8410d4c972f07bed5acbe59d935798c82fdbcd16

tools/intellij.lambda.testFramework/src/com/intellij/lambda/testFramework/junit/ExecuteInMonolithAndSplitMode.kt

@@ -5,7 +5,6 @@ import com.intellij.remoteDev.tests.modelGenerated.LambdaRdIdeType
 import com.intellij.remoteDev.tests.modelGenerated.LambdaRdIdeType.*
 import com.intellij.testFramework.junit5.impl.TestApplicationExtension
 import org.junit.jupiter.api.TestInstance
-import org.junit.jupiter.api.TestTemplate
 import org.junit.jupiter.api.extension.ExtendWith
 import java.lang.annotation.Inherited
 
@@ -14,7 +13,6 @@ import java.lang.annotation.Inherited
  */
 @Target(AnnotationTarget.ANNOTATION_CLASS, AnnotationTarget.CLASS, AnnotationTarget.FUNCTION)
 @Retention(AnnotationRetention.RUNTIME)
-@TestTemplate
 @Inherited
 @ExtendWith(TestApplicationExtension::class,
             MonolithAndSplitModeContextProvider::class,

tools/intellij.lambda.testFramework/src/com/intellij/lambda/testFramework/junit/InjectedLambda.kt

@@ -5,34 +5,254 @@ import com.intellij.openapi.diagnostic.logger
 import com.intellij.remoteDev.tests.LambdaFrontendContext
 import com.intellij.remoteDev.tests.impl.LambdaTestHost
 import com.intellij.remoteDev.tests.modelGenerated.LambdaRdKeyValueEntry
+import java.lang.invoke.MethodHandles
+import java.lang.invoke.MethodType
 import kotlin.reflect.full.createInstance
-import kotlin.reflect.full.declaredMembers
 
 private val logger = logger<InjectedLambda>()
 
+// TODO: Looks like JUnit test discovery and initialization should be reused here (probably from the monorepo source code)
 class InjectedLambda(frontendIdeContext: LambdaFrontendContext, plugin: PluginModuleDescriptor)
   : LambdaTestHost.Companion.NamedLambda<LambdaFrontendContext>(frontendIdeContext, plugin) {
 
-  // TODO: Looks like JUnit test discovery and initialization should be reused here (probably from the monorepo source code)
   override suspend fun LambdaFrontendContext.lambda(args: List<LambdaRdKeyValueEntry>): Any? {
     val className: String = args.singleOrNull { it.key == "testClass" }?.value
                             ?: error("Test class either not specified or specified multiple times. Args: $args")
     val methodName: String = args.singleOrNull { it.key == "testMethod" }?.value
                              ?: error("Test method either not specified or specified multiple times. Args: $args")
 
-    val testClass = Class.forName(className, true, plugin.pluginClassLoader).kotlin
+    // Verify method exists by reading bytecode directly
+    val methodInfo = findMethodInBytecode(className, methodName, plugin.pluginClassLoader!!)
+                     ?: error("Test method '$methodName' not found in test class '$className'")
 
-    val method = testClass.declaredMembers.singleOrNull { it.name == methodName }
-                 ?: error("Test method '$methodName' not found in test class '$className'")
-
-
-    val testContainer = testClass.createInstance()
+    val testClass = Class.forName(className, true, plugin.pluginClassLoader)
+    val testContainer = testClass.kotlin.createInstance()
     val rawArgs = argumentsFromString(args.singleOrNull { it.key == "methodArguments" }?.value ?: "")
-    val args: List<Any> = if (rawArgs.size == 1 && rawArgs.single() == "") listOf() else rawArgs
+    val methodArgs: List<Any> = if (rawArgs.size == 1 && rawArgs.single() == "") listOf() else rawArgs
 
-    logger.info("Starting test $className#$methodName inside ${lambdaIdeContext::class.simpleName} with args: $args")
+    logger.info("Starting test $className#$methodName inside ${lambdaIdeContext::class.simpleName} with args: $methodArgs")
 
-    return if (args.isNotEmpty()) method.call(testContainer, *args.toTypedArray())
-    else method.call(testContainer)
+    // Use MethodHandles to invoke without loading all methods
+    val lookup = MethodHandles.lookup()
+    val methodHandle = try {
+      if (methodArgs.isEmpty()) {
+        lookup.findVirtual(testClass, methodName, MethodType.methodType(Void.TYPE))
+      } else {
+        // For methods with parameters, we need to construct the method type
+        // This is a simplified version - you might need to handle parameter types
+        lookup.findVirtual(testClass, methodName, MethodType.methodType(Any::class.java, Array<Any>::class.java))
+      }
+    } catch (e: NoSuchMethodException) {
+      // Fallback: try to find the method by iterating one by one
+      findMethodSafely(testClass, methodName)?.let { method ->
+        method.isAccessible = true
+        return if (methodArgs.isNotEmpty()) method.invoke(testContainer, *methodArgs.toTypedArray())
+        else method.invoke(testContainer)
+      } ?: error("Method $methodName found in bytecode but cannot be invoked")
+    }
+
+    return if (methodArgs.isNotEmpty()) {
+      methodHandle.invokeWithArguments(testContainer, *methodArgs.toTypedArray())
+    } else {
+      methodHandle.invoke(testContainer)
+    }
+  }
+
+  /**
+   * Safely finds a method by trying to load methods one by one,
+   * catching exceptions for methods that cannot be loaded.
+   */
+  private fun findMethodSafely(clazz: Class<*>, methodName: String): java.lang.reflect.Method? {
+    // Try to get methods one by one using getDeclaredMethod
+    // But we don't know the parameter types, so we need to iterate differently
+    
+    // Use a safer approach: load the bytecode descriptor and try to match
+    try {
+      // Try no-arg method first (most common for test methods)
+      return clazz.getDeclaredMethod(methodName)
+    } catch (e: NoSuchMethodException) {
+      // Method has parameters, we need more sophisticated approach
+      logger.info("Method $methodName requires parameters or doesn't exist as no-arg")
+    } catch (e: NoClassDefFoundError) {
+      logger.warn("Cannot load method $methodName due to missing class: ${e.message}")
+    }
+    
+    // Last resort: try to iterate through declared methods with exception handling
+    val methods = mutableListOf<java.lang.reflect.Method>()
+    var index = 0
+    while (true) {
+      try {
+        // This is a hack: try to access methods by reflection on the internal array
+        val method = clazz.declaredMethods.getOrNull(index) ?: break
+        if (method.name == methodName) {
+          return method
+        }
+        index++
+      } catch (e: Throwable) {
+        // Skip methods that cause loading errors
+        logger.warn("Skipping method at index $index due to: ${e.message}")
+        index++
+        if (index > 100) break // Safety limit
+      }
+    }
+    
+    return null
+  }
+
+  /**
+   * Checks if a method exists in a class by reading the bytecode directly
+   * and returns method information.
+   */
+  private data class MethodInfo(val name: String, val paramCount: Int)
+
+  private fun findMethodInBytecode(className: String, methodName: String, classLoader: ClassLoader): MethodInfo? {
+    val classFileName = "${className.replace('.', '/')}.class"
+    
+    try {
+      val classBytes = classLoader.getResourceAsStream(classFileName)?.use { it.readBytes() }
+                       ?: run {
+                         logger.warn("Cannot find class file: $classFileName")
+                         return null
+                       }
+
+      val methods = parseMethodInfo(classBytes)
+      logger.info("Found methods in $className: ${methods.map { it.name }.joinToString()}")
+      
+      return methods.firstOrNull { it.name == methodName }
+    } catch (e: Exception) {
+      logger.error("Failed to read bytecode for $className", e)
+      return null
+    }
+  }
+
+  private fun parseMethodInfo(classBytes: ByteArray): List<MethodInfo> {
+    val methods = mutableListOf<MethodInfo>()
+    
+    try {
+      var offset = 0
+      
+      // Skip magic number (4 bytes) and version (4 bytes)
+      offset += 8
+      
+      // Read constant pool
+      val constantPoolCount = readUShort(classBytes, offset)
+      offset += 2
+      
+      val constantPool = Array<String?>(constantPoolCount) { null }
+      
+      // Parse constant pool
+      var i = 1
+      while (i < constantPoolCount) {
+        val tag = classBytes[offset].toInt() and 0xFF
+        offset++
+        
+        when (tag) {
+          1 -> { // CONSTANT_Utf8
+            val length = readUShort(classBytes, offset)
+            offset += 2
+            val str = String(classBytes, offset, length, Charsets.UTF_8)
+            constantPool[i] = str
+            offset += length
+          }
+          7, 8, 16, 19, 20 -> offset += 2
+          15 -> offset += 3
+          3, 4, 9, 10, 11, 12, 17, 18 -> offset += 4
+          5, 6 -> {
+            offset += 8
+            i++
+          }
+          else -> {
+            logger.warn("Unknown constant pool tag: $tag")
+            return emptyList()
+          }
+        }
+        i++
+      }
+      
+      // Skip access flags, this class, super class
+      offset += 6
+      
+      // Skip interfaces
+      val interfacesCount = readUShort(classBytes, offset)
+      offset += 2 + interfacesCount * 2
+      
+      // Skip fields
+      val fieldsCount = readUShort(classBytes, offset)
+      offset += 2
+      for (j in 0 until fieldsCount) {
+        offset += 6
+        offset = skipAttributes(classBytes, offset)
+      }
+      
+      // Read methods
+      val methodsCount = readUShort(classBytes, offset)
+      offset += 2
+      
+      for (j in 0 until methodsCount) {
+        offset += 2 // access_flags
+        val nameIndex = readUShort(classBytes, offset)
+        offset += 2
+        val descriptorIndex = readUShort(classBytes, offset)
+        offset += 2
+        
+        val name = constantPool[nameIndex]
+        val descriptor = constantPool[descriptorIndex]
+        
+        if (name != null && descriptor != null) {
+          val paramCount = countParameters(descriptor)
+          methods.add(MethodInfo(name, paramCount))
+        }
+        
+        offset = skipAttributes(classBytes, offset)
+      }
+      
+    } catch (e: Exception) {
+      logger.error("Error parsing class bytecode", e)
+    }
+    
+    return methods
+  }
+
+  private fun countParameters(descriptor: String): Int {
+    // Simple parameter counting from method descriptor
+    // e.g., "()V" = 0 params, "(I)V" = 1 param, "(ILjava/lang/String;)V" = 2 params
+    var count = 0
+    var inClass = false
+    for (c in descriptor.substringAfter('(').substringBefore(')')) {
+      when {
+        c == 'L' -> inClass = true
+        c == ';' -> {
+          inClass = false
+          count++
+        }
+        !inClass && c in "BCDFIJSZ" -> count++
+      }
+    }
+    return count
+  }
+
+  private fun readUShort(bytes: ByteArray, offset: Int): Int {
+    return ((bytes[offset].toInt() and 0xFF) shl 8) or (bytes[offset + 1].toInt() and 0xFF)
+  }
+
+  private fun skipAttributes(bytes: ByteArray, startOffset: Int): Int {
+    var offset = startOffset
+    val attributesCount = readUShort(bytes, offset)
+    offset += 2
+    
+    for (i in 0 until attributesCount) {
+      offset += 2
+      val attributeLength = readInt(bytes, offset)
+      offset += 4 + attributeLength
+    }
+    
+    return offset
+  }
+
+  private fun readInt(bytes: ByteArray, offset: Int): Int {
+    return ((bytes[offset].toInt() and 0xFF) shl 24) or
+           ((bytes[offset + 1].toInt() and 0xFF) shl 16) or
+           ((bytes[offset + 2].toInt() and 0xFF) shl 8) or
+           (bytes[offset + 3].toInt() and 0xFF)
   }
 }

tools/intellij.lambda.testFramework/src/com/intellij/lambda/testFramework/junit/MonolithAndSplitModeContextProvider.kt

@@ -1,5 +1,6 @@
 package com.intellij.lambda.testFramework.junit
 
+import com.intellij.lambda.testFramework.utils.BackgroundRunWithLambda
 import com.intellij.openapi.diagnostic.currentClassLogger
 import com.intellij.remoteDev.tests.modelGenerated.LambdaRdIdeType
 import org.junit.jupiter.api.extension.*
@@ -9,12 +10,11 @@ import org.junit.jupiter.params.provider.ArgumentsSource
 import org.junit.platform.commons.util.AnnotationUtils
 import java.util.*
 import java.util.function.Function
-import java.util.function.Supplier
 import java.util.stream.Stream
 
 // For the already existing tests that use test parametrization (and rewriting them to use @CartesianTest isn't desirable)
 class MonolithAndSplitModeContextProvider : TestTemplateInvocationContextProvider {
-  override fun supportsTestTemplate(context: ExtensionContext): Boolean = true
+  override fun supportsTestTemplate(context: ExtensionContext): Boolean = getModesToRun(context).isNotEmpty()
 
   private fun getModesToRun(context: ExtensionContext): List<LambdaRdIdeType> {
     val methodAnnotation = AnnotationUtils.findAnnotation(context.testMethod, ExecuteInMonolithAndSplitMode::class.java)
@@ -28,7 +28,7 @@ class MonolithAndSplitModeContextProvider : TestTemplateInvocationContextProvide
       return classAnnotation.get().mode.toList()
     }
 
-    throw IllegalStateException("The test should have ${ExecuteInMonolithAndSplitMode::javaClass.name} annotation")
+    throw IllegalStateException("The test is expected to have ${ExecuteInMonolithAndSplitMode::javaClass.name} annotation")
   }
 
   override fun provideTestTemplateInvocationContexts(context: ExtensionContext): Stream<TestTemplateInvocationContext> {
@@ -36,7 +36,8 @@ class MonolithAndSplitModeContextProvider : TestTemplateInvocationContextProvide
 
     currentClassLogger().info("Test will be run in modes: $modesToRun")
 
-    return getArgumentsProvider(context).map(Function { provider: ArgumentsProvider ->
+    // Test with parameters in the method signature
+    val parametrizedRuns = getArgumentsProvider(context).map(Function { provider: ArgumentsProvider ->
       try {
         // PARAMETERIZED CASE: Generate tests for each argument in each LambdaRdIdeType
         provider.provideArguments(context).flatMap { args: Arguments ->
@@ -46,10 +47,11 @@ class MonolithAndSplitModeContextProvider : TestTemplateInvocationContextProvide
       catch (e: Exception) {
         throw RuntimeException(e)
       }
-    }).orElseGet(Supplier { // SIMPLE CASE: Generate one test for each LambdaRdIdeType
-      modesToRun.stream().map { mode -> createInvocationContext(mode, arrayOf()) }
-    }
-    )
+    })
+    if (parametrizedRuns.isPresent) return parametrizedRuns.get()
+
+    // SIMPLE CASE - test without parameters: Generate one test for each LambdaRdIdeType
+    return modesToRun.stream().map { mode -> createInvocationContext(mode, emptyArray()) }
   }
 
   private fun createInvocationContext(mode: LambdaRdIdeType, args: Array<Any>): TestTemplateInvocationContext {
@@ -63,15 +65,23 @@ class MonolithAndSplitModeContextProvider : TestTemplateInvocationContextProvide
         return mutableListOf(object : ParameterResolver {
           override fun supportsParameter(paramCtx: ParameterContext, extCtx: ExtensionContext): Boolean {
             val type = paramCtx.parameter.type
-            return type == mode::class.java || (paramCtx.index > 0 && paramCtx.index - 1 < args.size)
+            return type.isAssignableFrom(BackgroundRunWithLambda::class.java) ||
+                   type.isAssignableFrom(mode::class.java) ||
+                       (paramCtx.index > 0 && paramCtx.index - 1 < args.size)
           }
 
           override fun resolveParameter(paramCtx: ParameterContext, extCtx: ExtensionContext): Any {
-            if (paramCtx.parameter.type == mode::class.java) {
-              return mode
+            return when {
+              paramCtx.parameter.type.isAssignableFrom(mode::class.java) -> mode
+              paramCtx.parameter.type.isAssignableFrom(BackgroundRunWithLambda::class.java) -> {
+                // Assuming BackgroundRunWithLambda requires specific instantiation logic
+                MonolithAndSplitModeIdeInstanceInitializer.ideBackgroundRun
+              }
+              else -> {
+                // The first parameter is the LambdaRdIdeType, so offset argument index by 1
+                args[paramCtx.index - 1]
+              }
             }
-            // The first parameter is the LambdaRdIdeType, so offset argument index by 1
-            return args[paramCtx.index - 1]
           }
         })
       }

tools/intellij.lambda.testFramework/src/com/intellij/lambda/testFramework/junit/MonolithAndSplitModeIdeInstanceInitializer.kt

@@ -5,9 +5,10 @@ import com.intellij.ide.starter.models.TestCase
 import com.intellij.ide.starter.project.LocalProjectInfo
 import com.intellij.ide.starter.project.TestCaseTemplate
 import com.intellij.ide.starter.runner.Starter
+import com.intellij.ide.starter.utils.catchAll
 import com.intellij.lambda.testFramework.starter.newContextWithLambda
+import com.intellij.lambda.testFramework.utils.BackgroundRunWithLambda
 import com.intellij.lambda.testFramework.utils.IdeLambdaStarter
-import com.intellij.lambda.testFramework.utils.IdeLambdaStarter.BackgroundRunWithLambda
 import com.intellij.lambda.testFramework.utils.IdeLambdaStarter.runIdeWithLambda
 import com.intellij.openapi.application.PathManager
 import org.junit.jupiter.api.extension.AfterAllCallback
@@ -29,7 +30,7 @@ class MonolithAndSplitModeIdeInstanceInitializer : BeforeAllCallback, AfterAllCa
   }
 
   override fun afterAll(context: ExtensionContext) {
-    ideBackgroundRun.closeIdeAndWait()
+    catchAll { ideBackgroundRun.closeIdeAndWait() }
   }
 }
 

tools/intellij.lambda.testFramework/src/com/intellij/lambda/testFramework/junit/MonolithAndSplitModeInvocationInterceptor.kt

@@ -1,5 +1,6 @@
 package com.intellij.lambda.testFramework.junit
 
+import com.intellij.lambda.testFramework.utils.BackgroundRunWithLambda
 import kotlinx.coroutines.runBlocking
 import org.junit.jupiter.api.extension.ExtensionContext
 import org.junit.jupiter.api.extension.InvocationInterceptor
@@ -41,14 +42,21 @@ open class MonolithAndSplitModeInvocationInterceptor : InvocationInterceptor {
   }
 
   private fun <T : Any> intercept(invocation: InvocationInterceptor.Invocation<T>, invocationContext: ReflectiveInvocationContext<Method>): T {
+    if (invocationContext.arguments.any { it::class == BackgroundRunWithLambda::class }) {
+      System.err.println("Test ${invocationContext.executable.name} has ${BackgroundRunWithLambda::class.qualifiedName} parameter. Test is expected to use it directly.")
+      invocation.proceed()
+    }
+
+    @Suppress("RAW_RUN_BLOCKING")
     runBlocking {
-      // TODO: init background IDE run before all tests (and provide and option to start IDE for every test)
+      // TODO: https://youtrack.jetbrains.com/issue/AT-3386/Lambda-tests-Start-different-instances-of-IDE-for-each-mode
+      // TODO: provide and option to start IDE for every test
       MonolithAndSplitModeIdeInstanceInitializer.ideBackgroundRun.runLambda(InjectedLambda::class,
                                                                             params = mapOf(
-                                                       "testClass" to invocationContext.targetClass.name,
-                                                       "testMethod" to invocationContext.executable.name,
-                                                       "methodArguments" to argumentsToString(invocationContext.arguments)
-                                                     ))
+                                                                              "testClass" to invocationContext.targetClass.name,
+                                                                              "testMethod" to invocationContext.executable.name,
+                                                                              "methodArguments" to argumentsToString(invocationContext.arguments)
+                                                                            ))
     }
 
     invocation.skip()

tools/intellij.lambda.testFramework/src/com/intellij/lambda/testFramework/utils/BackgroundRunWithLambda.kt

@@ -0,0 +1,45 @@
+package com.intellij.lambda.testFramework.utils
+
+import com.intellij.ide.starter.driver.engine.BackgroundRun
+import com.intellij.ide.starter.driver.engine.IBackgroundRun
+import com.intellij.lambda.testFramework.utils.IdeLambdaStarter.toLambdaParams
+import com.intellij.remoteDev.tests.LambdaBackendContext
+import com.intellij.remoteDev.tests.LambdaFrontendContext
+import com.intellij.remoteDev.tests.LambdaIdeContext
+import com.intellij.remoteDev.tests.impl.LambdaTestHost
+import com.intellij.remoteDev.tests.impl.utils.SerializedLambda
+import com.intellij.remoteDev.tests.modelGenerated.LambdaRdSerializedLambda
+import com.intellij.remoteDev.tests.modelGenerated.LambdaRdTestActionParameters
+import com.intellij.remoteDev.tests.modelGenerated.LambdaRdTestSession
+import kotlin.reflect.KClass
+
+class BackgroundRunWithLambda(delegate: BackgroundRun, val rdSession: LambdaRdTestSession, val backendRdSession: LambdaRdTestSession?) : IBackgroundRun by delegate {
+  internal suspend fun LambdaRdTestSession.runLambda(namedLambdaClass: KClass<out LambdaTestHost.Companion.NamedLambda<*>>, params: Map<String, String> = emptyMap()) {
+    runLambda.startSuspending(rdSession.protocol!!.lifetime,
+                              LambdaRdTestActionParameters(namedLambdaClass.java.canonicalName, params.toLambdaParams()))
+  }
+
+  suspend fun runLambda(namedLambdaClass: KClass<out LambdaTestHost.Companion.NamedLambda<*>>, params: Map<String, String> = emptyMap()) {
+    return rdSession.runLambda(namedLambdaClass, params)
+  }
+
+  suspend fun runLambdaInBackend(namedLambdaClass: KClass<out LambdaTestHost.Companion.NamedLambda<*>>, params: Map<String, String> = emptyMap()) {
+    return (backendRdSession ?: rdSession).runLambda(namedLambdaClass, params)
+  }
+
+
+  suspend inline fun <T : LambdaIdeContext> LambdaRdTestSession.runSerializedLambda(crossinline lambda: suspend T.() -> Unit) {
+    val exec = SerializedLambda.Companion.fromLambdaWithCoroutineScope(lambda)
+    runSerializedLambda.startSuspending(rdSession.protocol!!.lifetime,
+                                        LambdaRdSerializedLambda(exec.clazzName, exec.methodName, exec.serializedDataBase64,
+                                                                 exec.classPath.map { it.canonicalPath }))
+  }
+
+  suspend inline fun runSerializedLambda(crossinline lambda: suspend LambdaFrontendContext.() -> Unit) {
+    return rdSession.runSerializedLambda(lambda)
+  }
+
+  suspend inline fun runSerializedLambdaInBackend(crossinline lambda: suspend LambdaBackendContext.() -> Unit) {
+    return (backendRdSession ?: rdSession).runSerializedLambda(lambda)
+  }
+}

tools/intellij.lambda.testFramework/src/com/intellij/lambda/testFramework/utils/IdeLambdaStarter.kt

@@ -1,22 +1,18 @@
 package com.intellij.lambda.testFramework.utils
 
 import com.intellij.ide.starter.driver.driver.remoteDev.RemDevDriverRunner
-import com.intellij.ide.starter.driver.engine.BackgroundRun
-import com.intellij.ide.starter.driver.engine.IBackgroundRun
 import com.intellij.ide.starter.driver.engine.LocalDriverRunner
 import com.intellij.ide.starter.ide.IDERemDevTestContext
 import com.intellij.ide.starter.ide.IDETestContext
 import com.intellij.ide.starter.runner.IDERunContext
 import com.intellij.ide.starter.runner.events.IdeAfterLaunchEvent
-import com.intellij.remoteDev.tests.LambdaBackendContext
-import com.intellij.remoteDev.tests.LambdaFrontendContext
-import com.intellij.remoteDev.tests.LambdaIdeContext
 import com.intellij.remoteDev.tests.LambdaTestsConstants
-import com.intellij.remoteDev.tests.impl.LambdaTestHost.Companion.NamedLambda
 import com.intellij.remoteDev.tests.impl.LambdaTestHost.Companion.TEST_MODULE_ID_PROPERTY_NAME
-import com.intellij.remoteDev.tests.impl.utils.SerializedLambda
-import com.intellij.remoteDev.tests.modelGenerated.*
+import com.intellij.remoteDev.tests.modelGenerated.LambdaRdIdeType
 import com.intellij.remoteDev.tests.modelGenerated.LambdaRdIdeType.*
+import com.intellij.remoteDev.tests.modelGenerated.LambdaRdKeyValueEntry
+import com.intellij.remoteDev.tests.modelGenerated.LambdaRdTestSession
+import com.intellij.remoteDev.tests.modelGenerated.lambdaTestModel
 import com.intellij.remoteDev.util.executeSyncNonNullable
 import com.intellij.tools.ide.starter.bus.EventsBus
 import com.intellij.util.io.copyRecursively
@@ -25,48 +21,16 @@ import com.jetbrains.rd.framework.*
 import com.jetbrains.rd.util.lifetime.EternalLifetime
 import com.jetbrains.rd.util.threading.SynchronousScheduler
 import kotlin.io.path.exists
-import kotlin.reflect.KClass
 import kotlin.time.Duration
 import kotlin.time.Duration.Companion.minutes
 
-
 object IdeLambdaStarter {
   const val LAMBDA_TEST_MODULE: String = "intellij.lambda.tests._test"
   const val ADDITIONAL_LAMBDA_TEST_PLUGIN: String = "intellij.additional.lambda.tests.plugin"
   private const val ADDITIONAL_LAMBDA_DIR_NAME = "additional-lambda-tests-plugin"
 
-
   internal fun Map<String, String>.toLambdaParams(): List<LambdaRdKeyValueEntry> = map { LambdaRdKeyValueEntry(it.key, it.value) }
 
-  class BackgroundRunWithLambda(delegate: BackgroundRun, val rdSession: LambdaRdTestSession, val backendRdSession: LambdaRdTestSession?) : IBackgroundRun by delegate {
-    internal suspend fun LambdaRdTestSession.runLambda(namedLambdaClass: KClass<out NamedLambda<*>>, params: Map<String, String> = emptyMap()) {
-      runLambda.startSuspending(rdSession.protocol!!.lifetime,
-                                LambdaRdTestActionParameters(namedLambdaClass.java.canonicalName, params.toLambdaParams()))
-    }
-
-    suspend fun runLambda(namedLambdaClass: KClass<out NamedLambda<*>>, params: Map<String, String> = emptyMap()) {
-      return rdSession.runLambda(namedLambdaClass, params)
-    }
-
-    suspend fun runLambdaInBackend(namedLambdaClass: KClass<out NamedLambda<*>>, params: Map<String, String> = emptyMap()) {
-      return (backendRdSession ?: rdSession).runLambda(namedLambdaClass, params)
-    }
-
-
-    suspend inline fun <T : LambdaIdeContext> LambdaRdTestSession.runSerializedLambda(crossinline lambda: suspend T.() -> Unit) {
-      val exec = SerializedLambda.fromLambdaWithCoroutineScope(lambda)
-      runSerializedLambda.startSuspending(rdSession.protocol!!.lifetime, LambdaRdSerializedLambda(exec.clazzName, exec.methodName, exec.serializedDataBase64, exec.classPath.map { it.canonicalPath }))
-    }
-
-    suspend inline fun runSerializedLambda(crossinline lambda: suspend LambdaFrontendContext.() -> Unit) {
-      return rdSession.runSerializedLambda(lambda)
-    }
-
-    suspend inline fun runSerializedLambdaInBackend(crossinline lambda: suspend LambdaBackendContext.() -> Unit) {
-      return (backendRdSession ?: rdSession).runSerializedLambda(lambda)
-    }
-  }
-
   fun IDETestContext.runIdeWithLambda(
     runTimeout: Duration = 10.minutes,
     launchName: String = "",

tools/intellij.lambda.tests/tests/com/intellij/lambda/tests/FrontAndBackExampleTest.kt

@@ -0,0 +1,24 @@
+package com.intellij.lambda.tests
+
+import com.intellij.lambda.testFramework.junit.ExecuteInMonolithAndSplitMode
+import com.intellij.lambda.testFramework.utils.BackgroundRunWithLambda
+import com.intellij.lambda.tests.SampleTests.Companion.HelloBackendOnlyLambda
+import com.intellij.lambda.tests.SampleTests.Companion.HelloFrontendOnlyLambda
+import com.intellij.openapi.application.ApplicationManager
+import kotlinx.coroutines.runBlocking
+import org.junit.jupiter.api.TestTemplate
+
+@ExecuteInMonolithAndSplitMode
+class FrontAndBackExampleTest {
+
+  @TestTemplate
+  fun testTemplateTest(ide: BackgroundRunWithLambda) = runBlocking {
+    ide.runLambdaInBackend(HelloBackendOnlyLambda::class)
+    ide.runLambda(HelloFrontendOnlyLambda::class)
+  }
+
+  @TestTemplate
+  fun simpleUnitTest() {
+    ApplicationManager.getApplication().invokeAndWait { println("Test template test : badums") }
+  }
+}

tools/intellij.lambda.tests/tests/com/intellij/lambda/tests/SampleTests.kt

@@ -1,12 +1,12 @@
 package com.intellij.lambda.tests
 
 import com.intellij.ide.plugins.PluginModuleDescriptor
-import com.intellij.ide.starter.junit5.RemoteDevRun
-import com.intellij.lambda.testFramework.junit.MonolithAndSplitModeIdeInstanceInitializer
+import com.intellij.lambda.testFramework.junit.ExecuteInMonolithAndSplitMode
 import com.intellij.lambda.testFramework.testApi.editor.openFile
 import com.intellij.lambda.testFramework.testApi.getProject
 import com.intellij.lambda.testFramework.testApi.getProjects
 import com.intellij.lambda.testFramework.testApi.waitForProject
+import com.intellij.lambda.testFramework.utils.BackgroundRunWithLambda
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.diagnostic.currentClassLogger
 import com.intellij.remoteDev.tests.LambdaBackendContext
@@ -14,46 +14,37 @@ import com.intellij.remoteDev.tests.LambdaFrontendContext
 import com.intellij.remoteDev.tests.impl.LambdaTestHost.Companion.NamedLambda
 import com.intellij.remoteDev.tests.modelGenerated.LambdaRdKeyValueEntry
 import kotlinx.coroutines.runBlocking
-import org.junit.jupiter.api.Test
-import org.junit.jupiter.api.TestInstance
-import org.junit.jupiter.api.extension.ExtendWith
+import org.junit.jupiter.api.TestTemplate
 import kotlin.time.Duration.Companion.seconds
 
-@ExtendWith(RemoteDevRun::class, MonolithAndSplitModeIdeInstanceInitializer::class)
-@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+@ExecuteInMonolithAndSplitMode
 class SampleTests {
-  @Test
-  fun `serialized test`() {
+  @TestTemplate
+  fun `serialized test`(ide: BackgroundRunWithLambda) = runBlocking {
     //works in both modes if
     // TestCases.IU.JavaTestProject is used and
     // headless is turned off for monolith
-    runBlocking {
-      MonolithAndSplitModeIdeInstanceInitializer.ideBackgroundRun.apply {
-        runSerializedLambdaInBackend {
-          waitForProject(20.seconds)
-        }
+    ide.apply {
+      runSerializedLambdaInBackend {
+        waitForProject(20.seconds)
+      }
 
-        runSerializedLambda {
-          Logger.getInstance("test").warn("Projects: " + getProjects().joinToString { it.name })
-        }
+      runSerializedLambda {
+        Logger.getInstance("test").warn("Projects: " + getProjects().joinToString { it.name })
+      }
 
-        runSerializedLambdaInBackend {
-          Logger.getInstance("test").warn("backend Projects: " + getProject())
-          openFile("src/FormattingExamplesExpected.java", waitForReadyState = false)
-        }
+      runSerializedLambdaInBackend {
+        Logger.getInstance("test").warn("backend Projects: " + getProject())
+        openFile("src/FormattingExamplesExpected.java", waitForReadyState = false)
       }
     }
+    Unit
   }
 
-  @Test
-  fun `named lambda test`() {
-    runBlocking {
-      MonolithAndSplitModeIdeInstanceInitializer.ideBackgroundRun.apply {
-        runLambdaInBackend(HelloBackendOnlyLambda::class)
-
-        runLambda(HelloFrontendOnlyLambda::class)
-      }
-    }
+  @TestTemplate
+  fun `named lambda test`(ide: BackgroundRunWithLambda) = runBlocking {
+    ide.runLambdaInBackend(HelloBackendOnlyLambda::class)
+    ide.runLambda(HelloFrontendOnlyLambda::class)
   }
 
   companion object {