diff --git a/idea/testData/multiModuleQuickFix/withTest/commonTest/test/test.kt b/idea/testData/multiModuleQuickFix/withTest/commonTest/test/test.kt
index c0a6ecdaf20..9b67c877675 100644
--- a/idea/testData/multiModuleQuickFix/withTest/commonTest/test/test.kt
+++ b/idea/testData/multiModuleQuickFix/withTest/commonTest/test/test.kt
@@ -1,4 +1,5 @@
 // "Create actual function for platform JVM" "true"
+// TEST
 
 package test
 
diff --git a/idea/testData/multiModuleQuickFix/withTest/commonTest/test/test.kt.after b/idea/testData/multiModuleQuickFix/withTest/commonTest/test/test.kt.after
index 32a7b3e96f0..dc1c25aaa8d 100644
--- a/idea/testData/multiModuleQuickFix/withTest/commonTest/test/test.kt.after
+++ b/idea/testData/multiModuleQuickFix/withTest/commonTest/test/test.kt.after
@@ -1,4 +1,5 @@
 // "Create actual function for platform JVM" "true"
+// TEST
 
 package test
 
diff --git a/idea/testData/multiModuleQuickFix/withTest/jvmTest/test/testHelper.kt b/idea/testData/multiModuleQuickFix/withTest/jvmTest/test/testHelper.kt
index 2b774b34d6a..2395e79111a 100644
--- a/idea/testData/multiModuleQuickFix/withTest/jvmTest/test/testHelper.kt
+++ b/idea/testData/multiModuleQuickFix/withTest/jvmTest/test/testHelper.kt
@@ -1,3 +1,4 @@
 // testHelper: to be implemented
+// TEST
 
 package test
\ No newline at end of file
diff --git a/idea/testData/multiModuleQuickFix/withTest/jvmTest/test/testHelper.kt.after b/idea/testData/multiModuleQuickFix/withTest/jvmTest/test/testHelper.kt.after
index 8b23024f8d1..eb8f1a80fd9 100644
--- a/idea/testData/multiModuleQuickFix/withTest/jvmTest/test/testHelper.kt.after
+++ b/idea/testData/multiModuleQuickFix/withTest/jvmTest/test/testHelper.kt.after
@@ -1,4 +1,5 @@
 // testHelper: to be implemented
+// TEST
 
 package test
 
diff --git a/idea/tests/org/jetbrains/kotlin/idea/quickfix/AbstractQuickFixMultiModuleTest.kt b/idea/tests/org/jetbrains/kotlin/idea/quickfix/AbstractQuickFixMultiModuleTest.kt
index e3591c12713..30ebf35702d 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/quickfix/AbstractQuickFixMultiModuleTest.kt
+++ b/idea/tests/org/jetbrains/kotlin/idea/quickfix/AbstractQuickFixMultiModuleTest.kt
@@ -27,6 +27,8 @@ import org.jetbrains.kotlin.idea.project.PluginJetFilesProvider
 import org.jetbrains.kotlin.idea.stubs.AbstractMultiModuleTest
 import org.jetbrains.kotlin.idea.test.DirectiveBasedActionUtils
 import org.jetbrains.kotlin.idea.test.PluginTestCaseBase
+import org.jetbrains.kotlin.idea.util.module
+import org.jetbrains.kotlin.idea.util.sourceRoot
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.test.InTextDirectivesUtils
 import org.jetbrains.kotlin.test.KotlinTestUtils
@@ -67,14 +69,21 @@ abstract class AbstractQuickFixMultiModuleTest : AbstractMultiModuleTest() {
                     val testDirectory = File(testDataPath)
                     val projectDirectory = File("$testDataPath${getTestName(true)}")
                     for (moduleDirectory in projectDirectory.listFiles()) {
+                        val dirName = moduleDirectory.name
                         for (file in moduleDirectory.walkTopDown()) {
                             if (!file.path.endsWith(".after")) continue
                             try {
                                 val packageName = file.readLines().find { it.startsWith("package") }?.substringAfter(" ") ?: "<root>"
                                 val editedFile = allFilesInProject.mapNotNull {
                                     val candidate = it.containingDirectory?.findFile(file.name.removeSuffix(".after")) as? KtFile
-                                    if (candidate?.packageFqName?.toString() == packageName) candidate else null
-                                }.single()
+                                    val isUnderTestRoot = candidate != null && "// TEST" in candidate.text
+                                    candidate?.takeIf {
+                                        it.packageFqName.toString() == packageName &&
+                                        it.module?.let { module ->
+                                            module.name + ("Test".takeIf { isUnderTestRoot } ?: "") == dirName
+                                        } == true
+                                    }
+                                }.singleOrNull() ?: error("Cannot find suitable candidate for file ${file.name} from directory $dirName")
                                 setActiveEditor(editedFile.findExistingEditor() ?: createEditor(editedFile.virtualFile))
                                 checkResultByFile(file.relativeTo(testDirectory).path)
                             }