[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,28 +64,32 @@ class CompilerConfigurationProviderImpl(
 
     override fun getKotlinCoreEnvironment(module: TestModule): KotlinCoreEnvironment {
         return cache.getOrPut(module) {
-            val platform = module.targetPlatform
+            createKotlinCoreEnvironment(module)
-            val configFiles = when {
-                platform.isJvm() -> EnvironmentConfigFiles.JVM_CONFIG_FILES
-                platform.isJs() -> EnvironmentConfigFiles.JS_CONFIG_FILES
-                platform.isNative() -> EnvironmentConfigFiles.NATIVE_CONFIG_FILES
-                // TODO: is it correct?
-                platform.isCommon() -> EnvironmentConfigFiles.METADATA_CONFIG_FILES
-                else -> error("Unknown platform: $platform")
-            }
-            val applicationEnvironment = KotlinCoreEnvironment.getOrCreateApplicationEnvironmentForTests(
-                ApplicationEnvironmentDisposer.ROOT_DISPOSABLE,
-                CompilerConfiguration()
-            )
-            val projectEnv = KotlinCoreEnvironment.ProjectEnvironment(testRootDisposable, applicationEnvironment)
-            KotlinCoreEnvironment.createForTests(
-                projectEnv,
-                createCompilerConfiguration(module, projectEnv.project),
-                configFiles
-            )
         }
     }
 
+    protected open fun createKotlinCoreEnvironment(module: TestModule): KotlinCoreEnvironment {
+        val platform = module.targetPlatform
+        val configFiles = when {
+            platform.isJvm() -> EnvironmentConfigFiles.JVM_CONFIG_FILES
+            platform.isJs() -> EnvironmentConfigFiles.JS_CONFIG_FILES
+            platform.isNative() -> EnvironmentConfigFiles.NATIVE_CONFIG_FILES
+            // TODO: is it correct?
+            platform.isCommon() -> EnvironmentConfigFiles.METADATA_CONFIG_FILES
+            else -> error("Unknown platform: $platform")
+        }
+        val applicationEnvironment = KotlinCoreEnvironment.getOrCreateApplicationEnvironmentForTests(
+            ApplicationEnvironmentDisposer.ROOT_DISPOSABLE,
+            CompilerConfiguration()
+        )
+        val projectEnv = KotlinCoreEnvironment.ProjectEnvironment(testRootDisposable, applicationEnvironment)
+        return KotlinCoreEnvironment.createForTests(
+            projectEnv,
+            createCompilerConfiguration(module, projectEnv.project),
+            configFiles
+        )
+    }
+
     private fun createCompilerConfiguration(module: TestModule, project: MockProject): CompilerConfiguration {
         val configuration = CompilerConfiguration()
         configuration[CommonConfigurationKeys.MODULE_NAME] = module.name