[Test] Add ability to use custom CompilerConfigurationProvider

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/builders/TestConfigurationBuilder.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.test.builders
 
+import com.intellij.openapi.Disposable
 import org.jetbrains.kotlin.test.Constructor
 import org.jetbrains.kotlin.test.TestConfiguration
 import org.jetbrains.kotlin.test.TestInfrastructureInternals
@@ -41,6 +42,8 @@ class TestConfigurationBuilder {
     private val configurationsByTestDataCondition: MutableList<Pair<Regex, TestConfigurationBuilder.() -> Unit>> = mutableListOf()
     private val additionalServices: MutableList<ServiceRegistrationData> = mutableListOf()
 
+    private var compilerConfigurationProvider: ((Disposable, List<EnvironmentConfigurator>) -> CompilerConfigurationProvider)? = null
+
     lateinit var testInfo: KotlinTestInfo
 
     inline fun <reified T : TestService> useAdditionalService(noinline serviceConstructor: (TestServices) -> T) {
@@ -123,6 +126,11 @@ class TestConfigurationBuilder {
         moduleStructureTransformers += transformers
     }
 
+    @TestInfrastructureInternals
+    fun useCustomCompilerConfigurationProvider(provider: (Disposable, List<EnvironmentConfigurator>) -> CompilerConfigurationProvider) {
+        compilerConfigurationProvider = provider
+    }
+
     fun useMetaTestConfigurators(vararg configurators: Constructor<MetaTestConfigurator>) {
         metaTestConfigurators += configurators
     }
@@ -158,6 +166,7 @@ class TestConfigurationBuilder {
             moduleStructureTransformers,
             metaTestConfigurators,
             afterAnalysisCheckers,
+            compilerConfigurationProvider,
             metaInfoHandlerEnabled,
             directives,
             defaultRegisteredDirectivesBuilder.build(),

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/impl/TestConfigurationImpl.kt

@@ -37,6 +37,8 @@ class TestConfigurationImpl(
     metaTestConfigurators: List<Constructor<MetaTestConfigurator>>,
     afterAnalysisCheckers: List<Constructor<AfterAnalysisChecker>>,
 
+    compilerConfigurationProvider: ((Disposable, List<EnvironmentConfigurator>) -> CompilerConfigurationProvider)?,
+
     override val metaInfoHandlerEnabled: Boolean,
 
     directives: List<DirectivesContainer>,
@@ -96,7 +98,9 @@ class TestConfigurationImpl(
             val sourceFileProvider = SourceFileProviderImpl(this, sourceFilePreprocessors)
             register(SourceFileProvider::class, sourceFileProvider)
 
-            val environmentProvider = CompilerConfigurationProviderImpl(rootDisposable, this@TestConfigurationImpl.environmentConfigurators)
+            val environmentProvider =
+                compilerConfigurationProvider?.invoke(rootDisposable, this@TestConfigurationImpl.environmentConfigurators)
+                    ?: CompilerConfigurationProviderImpl(rootDisposable, this@TestConfigurationImpl.environmentConfigurators)
             register(CompilerConfigurationProvider::class, environmentProvider)
 
             register(AssertionsService::class, assertions)

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/services/CompilerConfigurationProvider.kt

@@ -56,7 +56,7 @@ abstract class CompilerConfigurationProvider : TestService {
 
 val TestServices.compilerConfigurationProvider: CompilerConfigurationProvider by TestServices.testServiceAccessor()
 
-class CompilerConfigurationProviderImpl(
+open class CompilerConfigurationProviderImpl(
     override val testRootDisposable: Disposable,
     val configurators: List<EnvironmentConfigurator>
 ) : CompilerConfigurationProvider() {
@@ -64,6 +64,11 @@ class CompilerConfigurationProviderImpl(
 
     override fun getKotlinCoreEnvironment(module: TestModule): KotlinCoreEnvironment {
         return cache.getOrPut(module) {
+            createKotlinCoreEnvironment(module)
+        }
+    }
+
+    protected open fun createKotlinCoreEnvironment(module: TestModule): KotlinCoreEnvironment {
         val platform = module.targetPlatform
         val configFiles = when {
             platform.isJvm() -> EnvironmentConfigFiles.JVM_CONFIG_FILES
@@ -78,13 +83,12 @@ class CompilerConfigurationProviderImpl(
             CompilerConfiguration()
         )
         val projectEnv = KotlinCoreEnvironment.ProjectEnvironment(testRootDisposable, applicationEnvironment)
-            KotlinCoreEnvironment.createForTests(
+        return KotlinCoreEnvironment.createForTests(
             projectEnv,
             createCompilerConfiguration(module, projectEnv.project),
             configFiles
         )
     }
-    }
 
     private fun createCompilerConfiguration(module: TestModule, project: MockProject): CompilerConfiguration {
         val configuration = CompilerConfiguration()