Support expecting different test outputs based on configuration

kotlin-analysis-api/src/test/kotlin/com/google/devtools/ksp/processor/TestFrameworkExpectDifferentOutputProcessor.kt

@@ -0,0 +1,38 @@
+/*
+ * Copyright 2026 Google LLC
+ * Copyright 2010-2026 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ *
+ * 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 com.google.devtools.ksp.processor
+
+import com.google.devtools.ksp.processing.Resolver
+import com.google.devtools.ksp.symbol.KSAnnotated
+
+class TestFrameworkExpectDifferentOutputProcessor(override val enableNewFeatures: Boolean) : AbstractTestProcessor() {
+    private val results = mutableListOf<String>()
+    override fun toResult(): List<String> = results
+
+    override fun process(resolver: Resolver): List<KSAnnotated> {
+        results.add("This should be in both configurations")
+        if (enableNewFeatures) {
+            results.add("This different output should be expected when enableNewFeatures = $enableNewFeatures")
+        } else {
+            results.add("This should be expected when enableNewFeatures = $enableNewFeatures")
+        }
+        results.add("This also be in both configurations")
+        return emptyList()
+    }
+
+}

kotlin-analysis-api/src/test/kotlin/com/google/devtools/ksp/test/AbstractKSPTest.kt

@@ -23,6 +23,7 @@ import com.intellij.openapi.Disposable
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Disposer
 import com.intellij.testFramework.TestDataFile
+import com.intellij.util.containers.toMultiMap
 import org.jetbrains.kotlin.analysis.test.framework.services.TargetPlatformDirectives
 import org.jetbrains.kotlin.analysis.test.framework.services.TargetPlatformProviderForAnalysisApiTests
 import org.jetbrains.kotlin.cli.common.disposeRootInWriteAction
@@ -54,6 +55,8 @@ import org.jetbrains.kotlin.test.services.configuration.CommonEnvironmentConfigu
 import org.jetbrains.kotlin.test.services.configuration.JvmEnvironmentConfigurator
 import org.jetbrains.kotlin.test.services.impl.TemporaryDirectoryManagerImpl
 import org.jetbrains.kotlin.test.util.KtTestUtil
+import org.jetbrains.kotlin.utils.addToStdlib.flatGroupBy
+import org.jetbrains.kotlin.utils.addToStdlib.getOrPut
 import org.junit.jupiter.api.AfterEach
 import org.junit.jupiter.api.Assertions
 import org.junit.jupiter.api.BeforeEach
@@ -82,12 +85,25 @@ abstract class DisposableTest {
 
 abstract class AbstractKSPTest(frontend: FrontendKind<*>, val enableNewFeatures: Boolean) : DisposableTest() {
     companion object {
-        const val TEST_PROCESSOR = "// TEST PROCESSOR:"
-        const val PROCESSOR_INPUT = "// PROCESSOR INPUT:"
-        const val EXPECTED_RESULTS = "// EXPECTED:"
-        const val EXPECTED_RESULTS_END = "// END"
-        const val MODULE = "// MODULE:"
-        const val COMPILER_MODULE_NAME = "// COMPILER MODULE NAME:"
+        const val COMMENT_TOKEN = "//"
+        const val TEST_PROCESSOR = "$COMMENT_TOKEN TEST PROCESSOR:"
+        const val PROCESSOR_INPUT = "$COMMENT_TOKEN PROCESSOR INPUT:"
+        const val EXPECTED_RESULTS = "$COMMENT_TOKEN EXPECTED:"
+
+        /**
+         * A directive controlling the expected test output when [enableNewFeatures] is `false`. The test output is
+         * expected to include the content on this line (modulo the directive).
+         */
+        const val EXPECT_CURRENT = "$COMMENT_TOKEN EXPECT CURRENT:"
+
+        /**
+         * A directive controlling the expected test output when [enableNewFeatures] is `true`. The test output is
+         * expected to include the content on this line (modulo the directive).
+         */
+        const val EXPECT_NEXT = "$COMMENT_TOKEN EXPECT NEXT:"
+        const val EXPECTED_RESULTS_END = "$COMMENT_TOKEN END"
+        const val MODULE = "$COMMENT_TOKEN MODULE:"
+        const val COMPILER_MODULE_NAME = "$COMMENT_TOKEN COMPILER MODULE NAME:"
     }
 
     init {
@@ -306,7 +322,9 @@ abstract class AbstractKSPTest(frontend: FrontendKind<*>, val enableNewFeatures:
         val processorClass = mkTestProcessorClass(parseTestProcessorName(fileContents))
         val testProcessor = mkProcessor(processorArguments, processorClass)
 
-        val expected = parseExpectedOutput(fileContents)
+        val expected = parseExpectedOutput(fileContents)[enableNewFeatures]
+            ?.joinToString("\n")
+            ?: ""
 
         val actual = {
             runTest(
@@ -334,14 +352,48 @@ abstract class AbstractKSPTest(frontend: FrontendKind<*>, val enableNewFeatures:
         ?.split(',')
         ?.map { it.trim() }
 
-    private fun parseExpectedOutput(fileContents: List<String>): String = fileContents
-        .dropWhile { !it.startsWith(EXPECTED_RESULTS) }
-        .drop(1)
-        .takeWhile { !it.startsWith(EXPECTED_RESULTS_END) }
-        .joinToString("\n") {
-            // Remove '// ' prefix
-            it.substring(3).trim()
+    /**
+     * Given the test file content, [parseExpectedOutput] returns a map of expected test results/output based
+     * on the [enableNewFeatures] feature toggle. Thus, given the feature toggle, the caller may index into the
+     * returned map to obtain the expected test results.
+     *
+     * [parseExpectedOutput] removes directives such as [EXPECT_CURRENT] and [EXPECT_NEXT] and removes dangling
+     * whitespace and comments. In other words, if `// MyExpectedOutput` is declared in the test file,
+     * the value `"MyExpectedOutput"` is in the returned list (for both configurations).
+     */
+    private fun parseExpectedOutput(fileContents: List<String>): Map<Boolean, List<String>> {
+        val rawExpectedOutput =
+            fileContents
+                .dropWhile { !it.startsWith(EXPECTED_RESULTS) }
+                .drop(1)
+                .takeWhile { !it.startsWith(EXPECTED_RESULTS_END) }
+
+        // Define simple aliases for readability
+        val newFeaturesDisabledConfiguration = false
+        val newFeaturesEnabledConfiguration = true
+
+        return buildMap<Boolean, MutableList<String>> {
+            rawExpectedOutput.forEach { line ->
+                when {
+                    line.startsWith(EXPECT_CURRENT) ->
+                        getOrPut(newFeaturesDisabledConfiguration, ::mutableListOf)
+                            .add(line.drop(EXPECT_CURRENT.length).trim())
+
+                    line.startsWith(EXPECT_NEXT) ->
+                        getOrPut(newFeaturesEnabledConfiguration, ::mutableListOf)
+                            .add(line.drop(EXPECT_NEXT.length).trim())
+
+                    else ->
+                        line.drop(COMMENT_TOKEN.length).trim().let {
+                            getOrPut(newFeaturesDisabledConfiguration, ::mutableListOf)
+                                .add(it)
+                            getOrPut(newFeaturesEnabledConfiguration, ::mutableListOf)
+                                .add(it)
+                        }
+                }
+            }
         }
+    }
 
     private fun mkProcessor(
         processorArguments: List<String>?,

kotlin-analysis-api/src/test/kotlin/com/google/devtools/ksp/test/KSPUnitTestSuite.kt

@@ -993,4 +993,35 @@ abstract class KSPUnitTestSuite(
     fun testPluginProblemReporter() {
         runThrowingTest("$AA_PATH/pluginProblemReporter.kt", expectedThrowableType = NullPointerException::class)
     }
+
+    @TestMetadata("expectDifferentOutput.kt")
+    @Test
+    fun testExpectDifferentOutput() {
+        runTest("$AA_PATH/expectDifferentOutput.kt")
+    }
+
+    @TestMetadata("expectDifferentOutputFailingOnCurrent.kt")
+    @Test
+    fun testExpectDifferentOutputFailingOnCurrent() {
+        // N.B.: This test is supposed to fail on one configuration.
+        // It asserts that the test output actually varies depending on the configuration.
+        if (enableNewFeatures) {
+            runTest("$AA_PATH/expectDifferentOutputFailingOnCurrent.kt")
+        } else {
+            runFailingTest("$AA_PATH/expectDifferentOutputFailingOnCurrent.kt")
+        }
+    }
+
+    @TestMetadata("expectDifferentOutputFailingOnNext.kt")
+    @Test
+    fun testExpectDifferentOutputFailingOnNext() {
+        // N.B.: This test is supposed to fail on one configuration.
+        // It asserts that the test output actually varies depending on the configuration.
+        if (enableNewFeatures) {
+            runFailingTest("$AA_PATH/expectDifferentOutputFailingOnNext.kt")
+        } else {
+            runTest("$AA_PATH/expectDifferentOutputFailingOnNext.kt")
+        }
+    }
+
 }

kotlin-analysis-api/testData/expectDifferentOutput.kt

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2026 Google LLC
+ * Copyright 2010-2026 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ *
+ * 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.
+ */
+
+// TEST PROCESSOR: TestFrameworkExpectDifferentOutputProcessor
+// EXPECTED:
+// This should be in both configurations
+// EXPECT CURRENT: This should be expected when enableNewFeatures = false
+// EXPECT NEXT: This different output should be expected when enableNewFeatures = true
+// This also be in both configurations
+// END
+
+// N.B.: Explicitly empty file. This test asserts that the test framework can change the expected test output with
+// the EXPECT CURRENT and EXPECT NEXT directives.

kotlin-analysis-api/testData/expectDifferentOutputFailingOnCurrent.kt

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2026 Google LLC
+ * Copyright 2010-2026 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ *
+ * 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.
+ */
+
+// TEST PROCESSOR: TestFrameworkExpectDifferentOutputProcessor
+// EXPECTED:
+// This should be in both configurations
+// EXPECT CURRENT: This should cause a failed test when enableNewFeatures = false
+// EXPECT NEXT: This different output should be expected when enableNewFeatures = true
+// This also be in both configurations
+// END
+
+// N.B.: Explicitly empty file. This test asserts that the test framework can change the expected test output with
+// the EXPECT CURRENT and EXPECT NEXT directives.

kotlin-analysis-api/testData/expectDifferentOutputFailingOnNext.kt

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2026 Google LLC
+ * Copyright 2010-2026 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ *
+ * 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.
+ */
+
+// TEST PROCESSOR: TestFrameworkExpectDifferentOutputProcessor
+// EXPECTED:
+// This should be in both configurations
+// EXPECT CURRENT: This should be expected when enableNewFeatures = false
+// EXPECT NEXT: This should cause a failure when enableNewFeatures = true
+// This also be in both configurations
+// END
+
+// N.B.: Explicitly empty file. This test asserts that the test framework can change the expected test output with
+// the EXPECT CURRENT and EXPECT NEXT directives.