diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileCrawler.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileCrawler.kt
index 47ecbd33d2f..b5a67642e44 100644
--- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileCrawler.kt
+++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileCrawler.kt
@@ -84,11 +84,16 @@ abstract class CodeWhispererFileCrawler : FileCrawler {
     }.orEmpty()
 
     override fun listCrossFileCandidate(target: PsiFile): List<VirtualFile> {
-        val targetFile = target.virtualFile
+        /**
+         * [PsiFile.getVirtualFile]
+         * PsiFile.virtualFile will return the virtual file, or null if the file exists only in memory
+         * If you want to get a non-null virtual file consider using FileViewProvider. getVirtualFile()
+         */
+        val targetVirtualFile = target.virtualFile ?: target.viewProvider.virtualFile
 
         val openedFiles = runReadAction {
             FileEditorManager.getInstance(target.project).openFiles.toList().filter {
-                it.name != target.virtualFile.name &&
+                it.name != targetVirtualFile.name &&
                     isSameDialect(it.extension) &&
                     !isTestFile(it, target.project)
             }
@@ -96,7 +101,7 @@ abstract class CodeWhispererFileCrawler : FileCrawler {
 
         val fileToFileDistanceList = runReadAction {
             openedFiles.map {
-                return@map it to CodeWhispererFileCrawler.getFileDistance(fileA = targetFile, fileB = it)
+                return@map it to CodeWhispererFileCrawler.getFileDistance(fileA = targetVirtualFile, fileB = it)
             }
         }
 
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt
index 62d5544bfd6..ab02be52873 100644
--- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt
+++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt
@@ -3,6 +3,7 @@
 
 package software.aws.toolkits.jetbrains.services.codewhisperer
 
+import com.intellij.openapi.application.readAction
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiFile
@@ -13,7 +14,6 @@ import com.intellij.testFramework.runInEdtAndWait
 import kotlinx.coroutines.test.runTest
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.Before
-import org.junit.Ignore
 import org.junit.Rule
 import org.junit.Test
 import org.mockito.kotlin.any
@@ -237,29 +237,51 @@ class CodeWhispererFileContextProviderTest {
         )
     }
 
-    /**
-     * - src/
-     *     - java/
-     *          - Main.java
-     *          - Util.java
-     *          - controllers/
-     *              -MyApiController.java
-     * - tst/
-     *     - java/
-     *          - MainTest.java
-     *
-     */
-    // TODO: fix this test, in test env, psiFile.virtualFile == null @psiGist.getFileData(psiFile) { psiFile -> ... }
-    @Ignore
     @Test
     fun `extractSupplementalFileContext from src file should extract src`() = runTest {
-        val psiFiles = setupFixture(fixture)
+        val file1 = "Human machine interface for lab abc computer applications"
+        val file2 = "A survey of user opinion of computer system response time"
+        val file3 = "The EPS user interface management system"
+        val file4 = "System and human system engineering testing of EPS"
+        val file5 = "Relation of user perceived response time to error measurement"
+        val file6 = "The generation of random binary unordered trees"
+        val file7 = "The intersection graph of paths in trees"
+        val file8 = "Graph minors IV Widths of trees and well quasi ordering"
+        val file9 = "Graph minors A survey"
+        val query = "The intersection of graph survey and trees"
+
+        val queryPsi = fixture.addFileToProject("Query.java", query)
+        val file1Psi = fixture.addFileToProject("File1.java", file1)
+        val file2Psi = fixture.addFileToProject("File2.java", file2)
+        val file3Psi = fixture.addFileToProject("File3.java", file3)
+        val file4Psi = fixture.addFileToProject("File4.java", file4)
+        val file5Psi = fixture.addFileToProject("File5.java", file5)
+        val file6Psi = fixture.addFileToProject("File6.java", file6)
+        val file7Psi = fixture.addFileToProject("File7.java", file7)
+        val file8Psi = fixture.addFileToProject("File8.java", file8)
+        val file9Psi = fixture.addFileToProject("File9.java", file9)
+
+        runInEdtAndWait {
+            fixture.openFileInEditor(file1Psi.viewProvider.virtualFile)
+            fixture.openFileInEditor(file2Psi.viewProvider.virtualFile)
+            fixture.openFileInEditor(file3Psi.viewProvider.virtualFile)
+            fixture.openFileInEditor(file4Psi.viewProvider.virtualFile)
+            fixture.openFileInEditor(file5Psi.viewProvider.virtualFile)
+            fixture.openFileInEditor(file6Psi.viewProvider.virtualFile)
+            fixture.openFileInEditor(file7Psi.viewProvider.virtualFile)
+            fixture.openFileInEditor(file8Psi.viewProvider.virtualFile)
+            fixture.openFileInEditor(file9Psi.viewProvider.virtualFile)
+        }
+
         sut = spy(sut)
 
-        val fileContext = sut.extractFileContext(fixture.editor, psiFiles[0])
+        val fileContext = readAction { sut.extractFileContext(fixture.editor, queryPsi) }
+        val supplementalContext = sut.extractSupplementalFileContext(queryPsi, fileContext, timeout = 50)
 
-        val supplementalContext = sut.extractSupplementalFileContext(psiFiles[0], fileContext, timeout = 50)
-        assertThat(supplementalContext?.contents).isNotNull.isNotEmpty
+        assertThat(supplementalContext?.contents)
+            .isNotNull
+            .isNotEmpty
+            .hasSize(3)
 
         verify(sut).extractSupplementalFileContextForSrc(any(), any())
         verify(sut, times(0)).extractSupplementalFileContextForTst(any(), any())