Generating unit test command not working

jetbrains/src/integrationTest/kotlin/com/sourcegraph/cody/util/BaseIntegrationTextFixture.kt

@@ -61,7 +61,7 @@ open class BaseIntegrationTextFixture(
     project = myFixture.project
     Disposer.register(myFixture.testRootDisposable) { shutdown() }
 
-    CodyEditorUtil.createFileOrScratchFromUntitled(
+    CodyEditorUtil.createFileOrUseExisting(
         project, ConfigUtil.getSettingsFile(project).toUri().toString(), codySettingsContent)
 
     initCredentialsAndAgent()

jetbrains/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentClient.kt

@@ -131,7 +131,7 @@ class CodyAgentClient(private val project: Project, private val webview: NativeW
       params: UntitledTextDocument
   ): CompletableFuture<ProtocolTextDocument?> {
     return acceptOnEventThreadAndGet {
-      val vf = CodyEditorUtil.createFileOrScratchFromUntitled(project, params.uri, params.content)
+      val vf = CodyEditorUtil.createFileOrUseExisting(project, params.uri, params.content)
       vf?.let { ProtocolTextDocumentExt.fromVirtualFile(it) }
     }
   }

jetbrains/src/main/kotlin/com/sourcegraph/cody/edit/EditService.kt

@@ -73,8 +73,7 @@ class EditService(val project: Project) {
         is CreateFileOperation -> {
           logger.info("Workspace edit operation created a file: ${op.uri}")
           val file =
-              CodyEditorUtil.createFileOrScratchFromUntitled(project, op.uri, content = "")
-                  ?: return false
+              CodyEditorUtil.createFileOrUseExisting(project, op.uri, content = "") ?: return false
           CodyEditorUtil.showDocument(project, file)
         }
         is RenameFileOperation -> {

jetbrains/src/main/kotlin/com/sourcegraph/common/CodyFileUri.kt

@@ -0,0 +1,72 @@
+package com.sourcegraph.common
+
+import com.intellij.openapi.vfs.VfsUtil
+import com.intellij.util.withPath
+import java.net.URI
+import java.net.URLDecoder
+import java.nio.file.InvalidPathException
+import java.nio.file.Path
+import java.nio.file.Paths
+import kotlin.io.path.Path
+import kotlin.io.path.toPath
+
+class CodyFileUri private constructor(val originalScheme: String, val uri: URI) {
+  val isUntitled: Boolean
+    get() = originalScheme == "untitled"
+
+  fun toPath(basePath: String?): Path {
+    var path = uri.toPath()
+    if (!path.isAbsolute && basePath != null) {
+      val fixedPath = Paths.get(path.toString().trimStart('/', '\\'))
+      path = Paths.get(basePath).resolve(fixedPath)
+    }
+    return path.normalize()
+  }
+
+  fun toPath(): Path = uri.toPath()
+
+  override fun toString() = uri.toString()
+
+  companion object {
+    private fun getPath(path: String): Path? {
+      return try {
+        Paths.get(path)
+      } catch (e: InvalidPathException) {
+        null
+      }
+    }
+
+    fun parse(input: String): CodyFileUri {
+      if (input.isEmpty()) throw IllegalArgumentException("input cannot be empty")
+
+      var processedInput = input
+      if (processedInput.contains("%")) {
+        processedInput = URLDecoder.decode(processedInput, "UTF-8")
+      }
+
+      var uri: URI
+      val scheme: String
+      val path = getPath(processedInput)
+      if (path != null) {
+        uri = path.toUri()
+        scheme = ""
+      } else {
+        val colonIndex = processedInput.indexOf(':')
+        scheme = processedInput.substring(0, colonIndex)
+        processedInput = "file:///${processedInput.substring(colonIndex + 1).trimStart('/')}"
+
+        uri =
+            VfsUtil.toUri(processedInput)
+                ?: throw IllegalArgumentException("input is not valid uri")
+        if (uri.path == null) {
+          uri = uri.withPath("/" + uri.schemeSpecificPart)
+        }
+        if (uri.path.contains("wsl.localhost")) {
+          uri = uri.withPath("////" + uri.path.trimStart('/'))
+        }
+      }
+
+      return CodyFileUri(scheme, uri)
+    }
+  }
+}

jetbrains/src/main/kotlin/com/sourcegraph/config/ConfigUtil.kt

@@ -198,7 +198,7 @@ object ConfigUtil {
   fun setCustomConfiguration(project: Project, customConfigContent: String): VirtualFile? {
     val config = ConfigFactory.parseString(customConfigContent).resolve()
     val content = config.root().render(renderOptions)
-    return CodyEditorUtil.createFileOrScratchFromUntitled(
+    return CodyEditorUtil.createFileOrUseExisting(
         project, getSettingsFile(project).toUri().toString(), content = content, overwrite = true)
         ?: run {
           logger.warn("Could not create settings file")

jetbrains/src/main/kotlin/com/sourcegraph/utils/CodyEditorUtil.kt

@@ -20,17 +20,15 @@ import com.intellij.openapi.project.ProjectManager
 import com.intellij.openapi.util.Key
 import com.intellij.openapi.util.TextRange
 import com.intellij.openapi.vfs.LocalFileSystem
-import com.intellij.openapi.vfs.VfsUtil
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.VirtualFileManager
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.intellij.util.io.createFile
 import com.sourcegraph.cody.agent.protocol_extensions.toOffsetRange
 import com.sourcegraph.cody.agent.protocol_generated.Range
+import com.sourcegraph.common.CodyFileUri
 import com.sourcegraph.config.ConfigUtil
 import com.sourcegraph.utils.ThreadingUtil.runInEdtAndGet
-import java.net.URISyntaxException
-import java.net.URLDecoder
 import kotlin.io.path.*
 
 object CodyEditorUtil {
@@ -177,62 +175,47 @@ object CodyEditorUtil {
     }
   }
 
-  @JvmStatic
-  fun fixUriString(uriString: String): String {
-    if (uriString.startsWith("untitled://")) {
-      // IntelliJ does not support in-memory files so we are using scratch files instead
-      return uriString.substringAfterLast(':').trimStart('/', '\\')
-    } else {
-      // Check `ProtocolTextDocumentExt.normalizeToVscUriFormat` for explanation
-      val patchedUri = uriString.replace("file://wsl.localhost/", "file:////wsl.localhost/")
-      return if (patchedUri.startsWith("file://")) patchedUri else "file://$patchedUri"
-    }
-  }
-
   fun findFileOrScratch(project: Project, uriString: String): VirtualFile? {
-    val fixedUri = fixUriString(uriString)
-    if (uriString.startsWith("untitled://")) {
+    val uri = CodyFileUri.parse(uriString)
+    if (uri.isUntitled) {
       return ScratchRootType.getInstance()
-          .findFile(project, fixedUri, ScratchFileService.Option.existing_only)
+          .findFile(project, uri.toString(), ScratchFileService.Option.existing_only)
     } else {
-      val uri = VfsUtil.toUri(fixedUri) ?: return null
-      return VirtualFileManager.getInstance().refreshAndFindFileByNioPath(uri.toPath())
+      val path = uri.toPath(project.basePath)
+      return VirtualFileManager.getInstance().refreshAndFindFileByNioPath(path)
     }
   }
 
-  fun createFileOrScratchFromUntitled(
+  fun createFileOrUseExisting(
       project: Project,
       uriString: String,
       content: String? = null,
       overwrite: Boolean = false
   ): VirtualFile? {
-    try {
-      val fileUri = VfsUtil.toUri(fixUriString(uriString)) ?: return null
-      if (overwrite || fileUri.toPath().notExists()) {
-        fileUri.toPath().parent?.createDirectories()
-        fileUri.toPath().deleteIfExists()
-        fileUri.toPath().createFile()
-        val vf = LocalFileSystem.getInstance().refreshAndFindFileByNioFile(fileUri.toPath())
-
-        content?.let {
-          WriteCommandAction.runWriteCommandAction(project) {
-            vf?.setBinaryContent(it.toByteArray())
-          }
-        }
+    val path = CodyFileUri.parse(uriString).toPath(project.basePath)
+    if (overwrite || path.notExists()) {
+      path.parent.createDirectories()
+      path.deleteIfExists()
+      path.createFile()
+      val vf = LocalFileSystem.getInstance().refreshAndFindFileByNioFile(path)
+
+      content?.let {
+        WriteCommandAction.runWriteCommandAction(project) { vf?.setBinaryContent(it.toByteArray()) }
       }
-
-      return LocalFileSystem.getInstance().refreshAndFindFileByNioFile(fileUri.toPath())
-    } catch (e: URISyntaxException) {
-      val fileName = uriString.substringAfterLast(':').trimStart('/', '\\')
-      val fileType = FileTypeRegistry.getInstance().getFileTypeByFileName(fileName)
-      val language = LanguageUtil.getFileTypeLanguage(fileType) ?: PlainTextLanguage.INSTANCE
-      return ScratchRootType.getInstance()
-          .createScratchFile(
-              project,
-              URLDecoder.decode(fileName, "UTF-8"),
-              language,
-              content ?: "",
-              ScratchFileService.Option.create_if_missing)
+      return vf
     }
+    return VirtualFileManager.getInstance().refreshAndFindFileByNioPath(path)
+  }
+
+  fun createScratchOrUseExisting(
+      project: Project,
+      fileName: String,
+      content: String? = null,
+  ): VirtualFile? {
+    val fileType = FileTypeRegistry.getInstance().getFileTypeByFileName(fileName)
+    val language = LanguageUtil.getFileTypeLanguage(fileType) ?: PlainTextLanguage.INSTANCE
+    return ScratchRootType.getInstance()
+        .createScratchFile(
+            project, fileName, language, content ?: "", ScratchFileService.Option.create_if_missing)
   }
 }

jetbrains/src/test/kotlin/com/sourcegraph/common/CodyFileUriTest.kt

@@ -0,0 +1,49 @@
+package com.sourcegraph.common
+
+import junit.framework.TestCase.assertEquals
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.Parameterized
+
+@RunWith(Parameterized::class)
+class CodyFileUriTest(private val uri: String, private val expectedUri: String) {
+
+  @Test
+  fun `can parse uri from agent`() {
+    val result = CodyFileUri.parse(uri)
+
+    assertEquals(expectedUri, result.toString())
+  }
+
+  companion object {
+    @JvmStatic
+    @Parameterized.Parameters
+    fun data(): Collection<Array<Any>> {
+      return listOf(
+          arrayOf(
+              "C:%5Cdev%5CJetbrainsTestsProjects%5Ckotlin%5CTestProject%5CValidate.kt",
+              "file:///C:/dev/JetbrainsTestsProjects/kotlin/TestProject/Validate.kt"),
+          arrayOf(
+              "C:\\dev\\JetbrainsTestsProjects\\kotlin\\TestProject\\Validate.kt",
+              "file:///C:/dev/JetbrainsTestsProjects/kotlin/TestProject/Validate.kt"),
+          arrayOf(
+              "file:///c%3A/dev/JetbrainsTestsProjects/kotlin/TestProject/src/Main.kt",
+              "file:///c:/dev/JetbrainsTestsProjects/kotlin/TestProject/src/Main.kt"),
+          arrayOf(
+              "file://c%3A/dev/JetbrainsTestsProjects/kotlin/TestProject/src/Main.kt",
+              "file:///c:/dev/JetbrainsTestsProjects/kotlin/TestProject/src/Main.kt"),
+          arrayOf("file:///c:/path/to/the%20file.txt", "file:///c:/path/to/the%20file.txt"),
+          arrayOf(
+              "untitled:/c%3A/dev/JetbrainsTestsProjects/kotlin/TestProject/src/test/kotlin/MainTest.kt",
+              "file:///c:/dev/JetbrainsTestsProjects/kotlin/TestProject/src/test/kotlin/MainTest.kt"),
+          arrayOf(
+              "untitled:7a8ee217-5b34-45c7-9a6e-eaf694ed3abc.kotlin",
+              "file:///7a8ee217-5b34-45c7-9a6e-eaf694ed3abc.kotlin"),
+          arrayOf(
+              "file:///Users/pk/Work/sourcegraph/cody/jetbrains/gradle.properties",
+              "file:///Users/pk/Work/sourcegraph/cody/jetbrains/gradle.properties"),
+          arrayOf("file://wsl.localhost/folder/file.cs", "file:////wsl.localhost/folder/file.cs"),
+      )
+    }
+  }
+}