From 5de2bebef76c63347200982a90b8f01b167eb7cd Mon Sep 17 00:00:00 2001
From: sksamuel <sam@sksamuel.com>
Date: Sun, 2 Jun 2024 19:46:40 -0500
Subject: [PATCH] Support navigation to tests

---
 .../kotlin/io/kotest/plugin/intellij/files.kt |  9 +--
 .../plugin/intellij/KotestTestLocator.kt      | 61 ++++++++++++-------
 2 files changed, 41 insertions(+), 29 deletions(-)

diff --git a/src/IC-242/kotlin/io/kotest/plugin/intellij/files.kt b/src/IC-242/kotlin/io/kotest/plugin/intellij/files.kt
index 3f047009..57a9807a 100644
--- a/src/IC-242/kotlin/io/kotest/plugin/intellij/files.kt
+++ b/src/IC-242/kotlin/io/kotest/plugin/intellij/files.kt
@@ -24,10 +24,7 @@ fun getLocationForFile(
    name: String,
    lineNumber: Int
 ): PsiLocation<PsiElement>? {
-   return FilenameIndex
-      .getVirtualFilesByName(name, scope)
-      .firstOrNull { it.isTestFile(project) }
-      ?.toPsiFile(project)
-      ?.elementAtLine(lineNumber)
-      ?.toPsiLocation()
+   val testFile = FilenameIndex.getVirtualFilesByName(name, scope).firstOrNull { it.isTestFile(project) } ?: return null
+   // element at line is 1 indexed, so we need to add one
+   return testFile.toPsiFile(project)?.elementAtLine(lineNumber + 1)?.toPsiLocation()
 }
diff --git a/src/main/kotlin/io/kotest/plugin/intellij/KotestTestLocator.kt b/src/main/kotlin/io/kotest/plugin/intellij/KotestTestLocator.kt
index b02ae763..a5339f49 100644
--- a/src/main/kotlin/io/kotest/plugin/intellij/KotestTestLocator.kt
+++ b/src/main/kotlin/io/kotest/plugin/intellij/KotestTestLocator.kt
@@ -25,34 +25,24 @@ import io.kotest.plugin.intellij.psi.toPsiLocation
  */
 class KotestTestLocator : SMTestLocator {
 
-   /**
-    * Returns the PSI file that contains the class indicated by the fully qualified name.
-    */
-   private fun loadPsiFile(fqn: String, project: Project, scope: GlobalSearchScope): PsiFile? {
-      val manager = PsiManager.getInstance(project)
-      // this is better than JavaPsiFacade as it will handle inner classes that use $
-      val psiClass = ClassUtil.findPsiClass(manager, fqn, null, true, scope)
-      val virtualFile = psiClass?.containingFile?.virtualFile
-      return virtualFile?.let { manager.findFile(it) }
-   }
-
-   private fun getLocationForFqn(
+   override fun getLocation(
+      protocol: String,
+      path: String,
       project: Project,
-      scope: GlobalSearchScope,
-      fqn: String,
-      lineNumber: Int
-   ): PsiLocation<PsiElement>? {
-      val psiFile = loadPsiFile(fqn, project, scope)
-      if (psiFile != null) {
-         val element = psiFile.elementAtLine(lineNumber) ?: psiFile
-         return element.toPsiLocation()
+      scope: GlobalSearchScope
+   ): List<Location<PsiElement>> {
+      return when (protocol) {
+         Constants().FileLocatorProtocol -> parseFile(project, scope, path)
+         Constants().ClassLocatorProtocol -> parseClass(project, scope, path)
+         Constants().OldLocatorProtocol -> parseClass(project, scope, path)
+         else -> emptyList()
       }
-      return null
    }
 
    override fun getLocation(
       protocol: String,
       path: String,
+      metainfo: String?,
       project: Project,
       scope: GlobalSearchScope
    ): List<Location<PsiElement>> {
@@ -64,6 +54,9 @@ class KotestTestLocator : SMTestLocator {
       }
    }
 
+   override fun getLocationCacheModificationTracker(project: Project): ModificationTracker =
+      ModificationTracker.EVER_CHANGED
+
    private fun parseFile(project: Project, scope: GlobalSearchScope, path: String): List<Location<PsiElement>> {
       val tokens = path.split(':')
       val ident = tokens[0]
@@ -78,8 +71,30 @@ class KotestTestLocator : SMTestLocator {
       return listOfNotNull(getLocationForFqn(project, scope, ident, lineNumber))
    }
 
-   override fun getLocationCacheModificationTracker(project: Project): ModificationTracker =
-      ModificationTracker.EVER_CHANGED
+   /**
+    * Returns the PSI file that contains the class indicated by the fully qualified name.
+    */
+   private fun loadPsiFile(fqn: String, project: Project, scope: GlobalSearchScope): PsiFile? {
+      val manager = PsiManager.getInstance(project)
+      // this is better than JavaPsiFacade as it will handle inner classes that use $
+      val psiClass = ClassUtil.findPsiClass(manager, fqn, null, true, scope)
+      val virtualFile = psiClass?.containingFile?.virtualFile
+      return virtualFile?.let { manager.findFile(it) }
+   }
+
+   private fun getLocationForFqn(
+      project: Project,
+      scope: GlobalSearchScope,
+      fqn: String,
+      lineNumber: Int
+   ): PsiLocation<PsiElement>? {
+      val psiFile = loadPsiFile(fqn, project, scope)
+      if (psiFile != null) {
+         val element = psiFile.elementAtLine(lineNumber) ?: psiFile
+         return element.toPsiLocation()
+      }
+      return null
+   }
 }