From b7cc781e9736e800587e590852e86da8b86c8368 Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Thu, 7 Jul 2022 14:10:18 +0300
Subject: [PATCH] [Test] Create KtFiles with proper virtual files

This is needed to properly report diagnostics in KAPT3 tests, which
  uses psiElement.containingFile as origin
---
 .../test/services/SourceFileProvider.kt       | 20 ++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/compiler/test-infrastructure/tests/org/jetbrains/kotlin/test/services/SourceFileProvider.kt b/compiler/test-infrastructure/tests/org/jetbrains/kotlin/test/services/SourceFileProvider.kt
index c2bd9f14332..dfca9041b5c 100644
--- a/compiler/test-infrastructure/tests/org/jetbrains/kotlin/test/services/SourceFileProvider.kt
+++ b/compiler/test-infrastructure/tests/org/jetbrains/kotlin/test/services/SourceFileProvider.kt
@@ -6,8 +6,14 @@
 package org.jetbrains.kotlin.test.services
 
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.text.StringUtilRt
+import com.intellij.openapi.vfs.StandardFileSystems
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiManager
+import com.intellij.testFramework.LightVirtualFile
 import org.jetbrains.kotlin.KtInMemoryTextSourceFile
 import org.jetbrains.kotlin.fir.lightTree.LightTree2Fir
+import org.jetbrains.kotlin.idea.KotlinLanguage
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.sourceFiles.LightTreeFile
 import org.jetbrains.kotlin.test.model.TestFile
@@ -86,9 +92,13 @@ class SourceFileProviderImpl(val testServices: TestServices, override val prepro
     }
 }
 
-fun SourceFileProvider.getKtFileForSourceFile(testFile: TestFile, project: Project): KtFile {
-// TODO
-//    return TestCheckerUtil.createCheckAndReturnPsiFile(
+fun SourceFileProvider.getKtFileForSourceFile(testFile: TestFile, project: Project, findViaVfs: Boolean = false): KtFile {
+    if (findViaVfs) {
+        val realFile = getRealFileForSourceFile(testFile)
+        StandardFileSystems.local().findFileByPath(realFile.path)
+            ?.let { PsiManager.getInstance(project).findFile(it) as? KtFile }
+            ?.let { return it }
+    }
     return KtTestUtil.createFile(
         testFile.name,
         getContentOfSourceFile(testFile),
@@ -96,10 +106,10 @@ fun SourceFileProvider.getKtFileForSourceFile(testFile: TestFile, project: Proje
     )
 }
 
-fun SourceFileProvider.getKtFilesForSourceFiles(testFiles: Collection<TestFile>, project: Project): Map<TestFile, KtFile> {
+fun SourceFileProvider.getKtFilesForSourceFiles(testFiles: Collection<TestFile>, project: Project, findViaVfs: Boolean = false): Map<TestFile, KtFile> {
     return testFiles.mapNotNull {
         if (!it.isKtFile) return@mapNotNull null
-        it to getKtFileForSourceFile(it, project)
+        it to getKtFileForSourceFile(it, project, findViaVfs)
     }.toMap()
 }