diff --git a/compiler/tests-common/tests/org/jetbrains/kotlin/checkers/BaseDiagnosticsTest.kt b/compiler/tests-common/tests/org/jetbrains/kotlin/checkers/BaseDiagnosticsTest.kt
index 629eae63d64..7bf93b2c559 100644
--- a/compiler/tests-common/tests/org/jetbrains/kotlin/checkers/BaseDiagnosticsTest.kt
+++ b/compiler/tests-common/tests/org/jetbrains/kotlin/checkers/BaseDiagnosticsTest.kt
@@ -91,7 +91,7 @@ abstract class BaseDiagnosticsTest : KotlinMultiFileTestWithJava<TestModule, Tes
         files: List<TestFile>
     ) {
         environment = createEnvironment(wholeFile)
-
+        setupEnvironment(environment)
         analyzeAndCheck(wholeFile, files)
     }
 
diff --git a/compiler/tests-common/tests/org/jetbrains/kotlin/checkers/KotlinMultiFileTestWithJava.java b/compiler/tests-common/tests/org/jetbrains/kotlin/checkers/KotlinMultiFileTestWithJava.java
index 3615e88c1eb..f542b1952df 100644
--- a/compiler/tests-common/tests/org/jetbrains/kotlin/checkers/KotlinMultiFileTestWithJava.java
+++ b/compiler/tests-common/tests/org/jetbrains/kotlin/checkers/KotlinMultiFileTestWithJava.java
@@ -93,6 +93,10 @@ public abstract class KotlinMultiFileTestWithJava<M extends KotlinBaseTest.TestM
 
     }
 
+    protected void setupEnvironment(@NotNull KotlinCoreEnvironment environment) {
+
+    }
+
     @NotNull
     protected ConfigurationKind getConfigurationKind() {
         return ConfigurationKind.JDK_ONLY;
diff --git a/compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirBaseDiagnosticsTest.kt b/compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirBaseDiagnosticsTest.kt
index 788f1059928..c565f65737a 100644
--- a/compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirBaseDiagnosticsTest.kt
+++ b/compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirBaseDiagnosticsTest.kt
@@ -54,12 +54,10 @@ abstract class AbstractFirBaseDiagnosticsTest : BaseDiagnosticsTest() {
         }
     }
 
-    override fun createEnvironment(file: File): KotlinCoreEnvironment {
-        return super.createEnvironment(file).apply {
-            Extensions.getArea(this.project)
-                .getExtensionPoint(PsiElementFinder.EP_NAME)
-                .unregisterExtension(JavaElementFinder::class.java)
-        }
+    override fun setupEnvironment(environment: KotlinCoreEnvironment) {
+        Extensions.getArea(environment.project)
+            .getExtensionPoint(PsiElementFinder.EP_NAME)
+            .unregisterExtension(JavaElementFinder::class.java)
     }
 
     open fun analyzeAndCheckUnhandled(testDataFile: File, files: List<TestFile>, useLightTree: Boolean = false) {
diff --git a/plugins/kotlin-serialization/kotlin-serialization-compiler/test/org/jetbrains/kotlinx/serialization/AbstractSerializationPluginDiagnosticTest.kt b/plugins/kotlin-serialization/kotlin-serialization-compiler/test/org/jetbrains/kotlinx/serialization/AbstractSerializationPluginDiagnosticTest.kt
index 75c643617e0..ccc4b7ba344 100644
--- a/plugins/kotlin-serialization/kotlin-serialization-compiler/test/org/jetbrains/kotlinx/serialization/AbstractSerializationPluginDiagnosticTest.kt
+++ b/plugins/kotlin-serialization/kotlin-serialization-compiler/test/org/jetbrains/kotlinx/serialization/AbstractSerializationPluginDiagnosticTest.kt
@@ -6,15 +6,15 @@
 package org.jetbrains.kotlinx.serialization
 
 import org.jetbrains.kotlin.checkers.AbstractDiagnosticsTest
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.cli.jvm.config.JvmClasspathRoot
 import org.jetbrains.kotlinx.serialization.compiler.extensions.SerializationComponentRegistrar
-import java.io.File
 
 abstract class AbstractSerializationPluginDiagnosticTest : AbstractDiagnosticsTest() {
     private val runtimeLibraryPath = getSerializationLibraryRuntimeJar()
 
-    override fun createEnvironment(file: File) = super.createEnvironment(file).apply {
-        SerializationComponentRegistrar.registerExtensions(this.project)
-        updateClasspath(listOf(JvmClasspathRoot(runtimeLibraryPath!!)))
+    override fun setupEnvironment(environment: KotlinCoreEnvironment) {
+        SerializationComponentRegistrar.registerExtensions(environment.project)
+        environment.updateClasspath(listOf(JvmClasspathRoot(runtimeLibraryPath!!)))
     }
 }
diff --git a/plugins/sam-with-receiver/sam-with-receiver-cli/test/org/jetbrains/kotlin/samWithReceiver/AbstractSamWithReceiverScriptTest.kt b/plugins/sam-with-receiver/sam-with-receiver-cli/test/org/jetbrains/kotlin/samWithReceiver/AbstractSamWithReceiverScriptTest.kt
index 4a129247cc9..918cb20788e 100644
--- a/plugins/sam-with-receiver/sam-with-receiver-cli/test/org/jetbrains/kotlin/samWithReceiver/AbstractSamWithReceiverScriptTest.kt
+++ b/plugins/sam-with-receiver/sam-with-receiver-cli/test/org/jetbrains/kotlin/samWithReceiver/AbstractSamWithReceiverScriptTest.kt
@@ -17,27 +17,24 @@
 package org.jetbrains.kotlin.samWithReceiver
 
 import org.jetbrains.kotlin.checkers.AbstractDiagnosticsTest
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
 import org.jetbrains.kotlin.scripting.configuration.ScriptingConfigurationKeys
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
 import org.jetbrains.kotlin.scripting.resolve.KotlinScriptDefinitionFromAnnotatedTemplate
-import java.io.File
 import kotlin.script.experimental.jvm.defaultJvmScriptingHostConfiguration
 import kotlin.script.extensions.SamWithReceiverAnnotations
 import kotlin.script.templates.ScriptTemplateDefinition
 
 abstract class AbstractSamWithReceiverScriptTest : AbstractDiagnosticsTest() {
-    private companion object {
-        private val TEST_ANNOTATIONS = emptyList<String>()
-    }
 
-    override fun createEnvironment(file: File) = super.createEnvironment(file).apply {
-        StorageComponentContainerContributor.registerExtension(project, CliSamWithReceiverComponentContributor(TEST_ANNOTATIONS))
+    override fun setupEnvironment(environment: KotlinCoreEnvironment) {
+        StorageComponentContainerContributor.registerExtension(environment.project, CliSamWithReceiverComponentContributor(emptyList()))
         val def = ScriptDefinition.FromLegacy(
             defaultJvmScriptingHostConfiguration,
             KotlinScriptDefinitionFromAnnotatedTemplate(ScriptForSamWithReceivers::class, emptyMap())
         )
-        configuration.add(ScriptingConfigurationKeys.SCRIPT_DEFINITIONS, def)
+        environment.configuration.add(ScriptingConfigurationKeys.SCRIPT_DEFINITIONS, def)
     }
 }
 
diff --git a/plugins/sam-with-receiver/sam-with-receiver-cli/test/org/jetbrains/kotlin/samWithReceiver/AbstractSamWithReceiverTest.kt b/plugins/sam-with-receiver/sam-with-receiver-cli/test/org/jetbrains/kotlin/samWithReceiver/AbstractSamWithReceiverTest.kt
index fe62f92b893..8ae4122978f 100644
--- a/plugins/sam-with-receiver/sam-with-receiver-cli/test/org/jetbrains/kotlin/samWithReceiver/AbstractSamWithReceiverTest.kt
+++ b/plugins/sam-with-receiver/sam-with-receiver-cli/test/org/jetbrains/kotlin/samWithReceiver/AbstractSamWithReceiverTest.kt
@@ -17,15 +17,18 @@
 package org.jetbrains.kotlin.samWithReceiver
 
 import org.jetbrains.kotlin.checkers.AbstractDiagnosticsTest
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
-import java.io.File
 
 abstract class AbstractSamWithReceiverTest : AbstractDiagnosticsTest() {
     private companion object {
         private val TEST_ANNOTATIONS = listOf("SamWithReceiver")
     }
 
-    override fun createEnvironment(file: File) = super.createEnvironment(file).apply {
-        StorageComponentContainerContributor.registerExtension(project, CliSamWithReceiverComponentContributor(TEST_ANNOTATIONS))
+    override fun setupEnvironment(environment: KotlinCoreEnvironment) {
+        StorageComponentContainerContributor.registerExtension(
+            environment.project,
+            CliSamWithReceiverComponentContributor(TEST_ANNOTATIONS)
+        )
     }
 }