diff --git a/compiler/tests-common/org/jetbrains/kotlin/codegen/AbstractBytecodeListingTest.kt b/compiler/tests-common/org/jetbrains/kotlin/codegen/AbstractBytecodeListingTest.kt
index 1f1719a95bf..9d8ac3be068 100644
--- a/compiler/tests-common/org/jetbrains/kotlin/codegen/AbstractBytecodeListingTest.kt
+++ b/compiler/tests-common/org/jetbrains/kotlin/codegen/AbstractBytecodeListingTest.kt
@@ -70,6 +70,8 @@ abstract class AbstractBytecodeListingTest : CodegenTestCase() {
                 testRootDisposable, configuration, EnvironmentConfigFiles.JVM_CONFIG_FILES
         )
 
+        setupEnvironment(myEnvironment)
+
         loadMultiFiles(files)
 
         val txtFile = getTextFile(wholeFile)
@@ -86,6 +88,8 @@ abstract class AbstractBytecodeListingTest : CodegenTestCase() {
         KotlinTestUtils.assertEqualsToFile(txtFile, generatedFiles)
     }
 
+    protected open fun setupEnvironment(environment: KotlinCoreEnvironment) {}
+
     private class TextCollectingVisitor : ClassVisitor(ASM5) {
         private class Declaration(val text: String, val annotations: MutableList<String> = arrayListOf())
 
diff --git a/compiler/tests-common/org/jetbrains/kotlin/codegen/AbstractLightAnalysisModeCodegenTest.kt b/compiler/tests-common/org/jetbrains/kotlin/codegen/AbstractLightAnalysisModeCodegenTest.kt
index 4630fdb8e20..f862541d8dc 100644
--- a/compiler/tests-common/org/jetbrains/kotlin/codegen/AbstractLightAnalysisModeCodegenTest.kt
+++ b/compiler/tests-common/org/jetbrains/kotlin/codegen/AbstractLightAnalysisModeCodegenTest.kt
@@ -16,15 +16,23 @@
 
 package org.jetbrains.kotlin.codegen
 
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
+import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisHandlerExtension
+import org.jetbrains.kotlin.resolve.jvm.extensions.PartialAnalysisHandlerExtension
 import java.io.File
 
 abstract class AbstractLightAnalysisModeCodegenTest : AbstractBytecodeListingTest() {
     override val classBuilderFactory: ClassBuilderFactory
-        get() = ClassBuilderFactories.TEST_LIGHT_CLASSES
+        get() = ClassBuilderFactories.TEST_KAPT3
+
+    override fun setupEnvironment(environment: KotlinCoreEnvironment) {
+        AnalysisHandlerExtension.registerExtension(environment.project, PartialAnalysisHandlerExtension())
+    }
+
 
     override fun getTextFile(ktFile: File): File {
         val boxTestsDir = File("compiler/testData/codegen/box")
         val outDir = File("compiler/testData/codegen/light-analysis", ktFile.toRelativeString(boxTestsDir)).parent
         return File(outDir, ktFile.nameWithoutExtension + ".txt")
     }
-}
+}
\ No newline at end of file