diff --git a/plugins/kapt3/src/org/jetbrains/kotlin/kapt3/Kapt3Extension.kt b/plugins/kapt3/src/org/jetbrains/kotlin/kapt3/Kapt3Extension.kt
index 9e8061115bc..269aa6999a8 100644
--- a/plugins/kapt3/src/org/jetbrains/kotlin/kapt3/Kapt3Extension.kt
+++ b/plugins/kapt3/src/org/jetbrains/kotlin/kapt3/Kapt3Extension.kt
@@ -32,7 +32,6 @@ import org.jetbrains.kotlin.kapt3.util.KaptLogger
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.BindingTrace
-import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisHandlerExtension
 import org.jetbrains.kotlin.resolve.jvm.extensions.PartialAnalysisHandlerExtension
 import java.io.File
 import java.net.URLClassLoader
@@ -40,8 +39,34 @@ import java.util.*
 import javax.annotation.processing.Processor
 import com.sun.tools.javac.util.List as JavacList
 
-class Kapt3Extension(
+class ClasspathBasedKapt3Extension(
         val annotationProcessingClasspath: List<File>,
+        javaSourceRoots: List<File>,
+        sourcesOutputDir: File,
+        classFilesOutputDir: File,
+        stubsOutputDir: File?,
+        options: Map<String, String>,
+        aptOnly: Boolean,
+        pluginInitializedTime: Long,
+        logger: KaptLogger
+) : AbstractKapt3Extension(annotationProcessingClasspath, javaSourceRoots, sourcesOutputDir,
+                           classFilesOutputDir, stubsOutputDir, options, aptOnly, pluginInitializedTime, logger) {
+    override fun loadProcessors(): List<Processor> {
+        val classLoader = URLClassLoader(annotationProcessingClasspath.map { it.toURI().toURL() }.toTypedArray())
+        val processors = ServiceLoader.load(Processor::class.java, classLoader).toList()
+
+        if (processors.isEmpty()) {
+            logger.info("No annotation processors available, aborting")
+        } else {
+            logger.info { "Annotation processors: " + processors.joinToString { it.javaClass.canonicalName } }
+        }
+
+        return processors
+    }
+}
+
+abstract class AbstractKapt3Extension(
+        val classpath: List<File>,
         val javaSourceRoots: List<File>,
         val sourcesOutputDir: File,
         val classFilesOutputDir: File,
@@ -78,7 +103,7 @@ class Kapt3Extension(
         logger.info { "Initial analysis took ${System.currentTimeMillis() - pluginInitializedTime} ms" }
         logger.info { "Kotlin files to compile: " + files.map { it.virtualFile?.name ?: "<in memory ${it.hashCode()}>" } }
 
-        val processors = loadProcessors(annotationProcessingClasspath)
+        val processors = loadProcessors()
         if (processors.isEmpty()) return if (aptOnly) doNotGenerateCode() else null
 
         val (kaptContext, generationState) = compileStubs(project, module, bindingTrace.bindingContext, files.toList())
@@ -90,7 +115,7 @@ class Kapt3Extension(
             val (annotationProcessingTime) = measureTimeMillis {
                 kaptContext.doAnnotationProcessing(
                         javaSourceFiles, processors,
-                        annotationProcessingClasspath, sourcesOutputDir, classFilesOutputDir, kotlinSourceStubs)
+                        classpath, sourcesOutputDir, classFilesOutputDir, kotlinSourceStubs)
             }
 
             logger.info { "Annotation processing took $annotationProcessingTime ms" }
@@ -180,18 +205,7 @@ class Kapt3Extension(
         return javaFilesFromJavaSourceRoots
     }
 
-    private fun loadProcessors(classpath: List<File>): List<Processor> {
-        val classLoader = URLClassLoader(classpath.map { it.toURI().toURL() }.toTypedArray())
-        val processors = ServiceLoader.load(Processor::class.java, classLoader).toList()
-
-        if (processors.isEmpty()) {
-            logger.info("No annotation processors available, aborting")
-        } else {
-            logger.info { "Annotation processors: " + processors.joinToString { it.javaClass.canonicalName } }
-        }
-
-        return processors
-    }
+    protected abstract fun loadProcessors(): List<Processor>
 }
 
 private inline fun <T> measureTimeMillis(block: () -> T) : Pair<Long, T> {
diff --git a/plugins/kapt3/src/org/jetbrains/kotlin/kapt3/Kapt3Plugin.kt b/plugins/kapt3/src/org/jetbrains/kotlin/kapt3/Kapt3Plugin.kt
index 5be6e2a6710..0f34498b23c 100644
--- a/plugins/kapt3/src/org/jetbrains/kotlin/kapt3/Kapt3Plugin.kt
+++ b/plugins/kapt3/src/org/jetbrains/kotlin/kapt3/Kapt3Plugin.kt
@@ -152,7 +152,7 @@ class Kapt3ComponentRegistrar : ComponentRegistrar {
             logger.info("Options: $apOptions")
         }
 
-        val kapt3AnalysisCompletedHandlerExtension = Kapt3Extension(
+        val kapt3AnalysisCompletedHandlerExtension = ClasspathBasedKapt3Extension(
                 classpath, javaSourceRoots, sourcesOutputDir, classFilesOutputDir, stubsOutputDir, apOptions,
                 isAptOnly, System.currentTimeMillis(), logger)
         AnalysisHandlerExtension.registerExtension(project, kapt3AnalysisCompletedHandlerExtension)
diff --git a/plugins/kapt3/test/org/jetbrains/kotlin/kapt3/test/AbstractKotlinKapt3Test.kt b/plugins/kapt3/test/org/jetbrains/kotlin/kapt3/test/AbstractKotlinKapt3Test.kt
index eadeef2134e..188d517d85b 100644
--- a/plugins/kapt3/test/org/jetbrains/kotlin/kapt3/test/AbstractKotlinKapt3Test.kt
+++ b/plugins/kapt3/test/org/jetbrains/kotlin/kapt3/test/AbstractKotlinKapt3Test.kt
@@ -36,7 +36,7 @@ import java.io.File
 import java.nio.file.Files
 
 abstract class AbstractKotlinKapt3Test : CodegenTestCase() {
-    protected companion object {
+    companion object {
         val FILE_SEPARATOR = "\n\n////////////////////\n\n"
     }
 
diff --git a/plugins/kapt3/test/org/jetbrains/kotlin/kapt3/test/AbstractKotlinKapt3TestWithExtension.kt b/plugins/kapt3/test/org/jetbrains/kotlin/kapt3/test/AbstractKotlinKapt3TestWithExtension.kt
new file mode 100644
index 00000000000..d3c1ba6ce04
--- /dev/null
+++ b/plugins/kapt3/test/org/jetbrains/kotlin/kapt3/test/AbstractKotlinKapt3TestWithExtension.kt
@@ -0,0 +1,76 @@
+/*
+ * Copyright 2010-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.kapt3.test
+
+import com.intellij.openapi.util.text.StringUtil
+import org.jetbrains.kotlin.codegen.CodegenTestCase
+import org.jetbrains.kotlin.codegen.CodegenTestUtil
+import org.jetbrains.kotlin.kapt3.AbstractKapt3Extension
+import org.jetbrains.kotlin.kapt3.Kapt3BuilderFactory
+import org.jetbrains.kotlin.kapt3.util.KaptLogger
+import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisHandlerExtension
+import org.jetbrains.kotlin.test.ConfigurationKind
+import org.jetbrains.kotlin.test.KotlinTestUtils
+import org.jetbrains.kotlin.test.util.trimTrailingWhitespacesAndAddNewlineAtEOF
+import org.junit.Test
+import java.io.File
+import java.nio.file.Files
+import javax.annotation.processing.Processor
+
+class KotlinKapt3TestWithExtension : CodegenTestCase() {
+    class Kapt3ExtensionForTests(
+            private val processors: List<Processor>,
+            javaSourceRoots: List<File>,
+            outputDir: File
+    ) : AbstractKapt3Extension(emptyList(), javaSourceRoots, outputDir, outputDir,
+                               null, emptyMap(), true, System.currentTimeMillis(), KaptLogger(true)) {
+        override fun loadProcessors() = processors
+    }
+
+    private fun getProcessor(): Processor = KaptRunnerTest.simpleProcessor()
+
+    @Test
+    fun testSimple() {
+        doTest("plugins/kapt3/testData/kotlinRunner/Simple.kt")
+    }
+
+    override fun doMultiFileTest(wholeFile: File, files: List<TestFile>, javaFilesDir: File?) {
+        val javaSources = javaFilesDir?.let { arrayOf(it) } ?: emptyArray()
+
+        val txtFile = File(wholeFile.parentFile, wholeFile.nameWithoutExtension + ".ext.txt")
+        val sourceOutputDir = Files.createTempDirectory("kaptRunner").toFile()
+
+        createEnvironmentWithMockJdkAndIdeaAnnotations(ConfigurationKind.ALL, *javaSources)
+        val kapt3Extension = Kapt3ExtensionForTests(listOf(getProcessor()), javaSources.toList(), sourceOutputDir)
+        AnalysisHandlerExtension.registerExtension(myEnvironment.project, kapt3Extension)
+
+        try {
+            loadMultiFiles(files)
+
+            val classBuilderFactory = Kapt3BuilderFactory()
+            CodegenTestUtil.generateFiles(myEnvironment, myFiles, classBuilderFactory)
+
+            val javaFiles = sourceOutputDir.walkTopDown().filter { it.isFile && it.extension == "java" }
+            val actualRaw = javaFiles.sortedBy { it.name }.joinToString(AbstractKotlinKapt3Test.FILE_SEPARATOR) { it.name + ":\n\n" + it.readText() }
+            val actual = StringUtil.convertLineSeparators(actualRaw.trim({ it <= ' ' })).trimTrailingWhitespacesAndAddNewlineAtEOF()
+            KotlinTestUtils.assertEqualsToFile(txtFile, actual)
+        } finally {
+            sourceOutputDir.deleteRecursively()
+        }
+    }
+
+}
\ No newline at end of file
diff --git a/plugins/kapt3/testData/kotlinRunner/Simple.ext.txt b/plugins/kapt3/testData/kotlinRunner/Simple.ext.txt
new file mode 100644
index 00000000000..214dd6af195
--- /dev/null
+++ b/plugins/kapt3/testData/kotlinRunner/Simple.ext.txt
@@ -0,0 +1,4 @@
+MyMethodMyAnnotation.java:
+
+package generated;
+class MyMethodMyAnnotation {}
diff --git a/plugins/kapt3/testData/kotlinRunner/Simple.kt b/plugins/kapt3/testData/kotlinRunner/Simple.kt
index fb348eb9478..4dabfb6feef 100644
--- a/plugins/kapt3/testData/kotlinRunner/Simple.kt
+++ b/plugins/kapt3/testData/kotlinRunner/Simple.kt
@@ -5,6 +5,10 @@ internal class Simple {
     fun myMethod() {
         // do nothing
     }
+
+    fun heavyMethod(): Int {
+        return if (true) 5 else 6
+    }
 }
 
 internal annotation class MyAnnotation