diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererUTGChatManager.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererUTGChatManager.kt
index c25462b769a..e1419c4c107 100644
--- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererUTGChatManager.kt
+++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererUTGChatManager.kt
@@ -524,6 +524,8 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
             try {
                 launchTestGenFlow(prompt, codeTestChatHelper, previousIterationContext, selectionRange)
             } catch (e: Exception) {
+                // reset number of unitTestGenerated to null
+                session.numberOfUnitTestCasesGenerated = null
                 // Add an answer for displaying error message
                 val errorMessage = when {
                     e is CodeTestException &&
diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/controller/CodeTestChatController.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/controller/CodeTestChatController.kt
index 712167aa7b2..9abaf0e5d27 100644
--- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/controller/CodeTestChatController.kt
+++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/controller/CodeTestChatController.kt
@@ -9,7 +9,6 @@ import com.intellij.diff.DiffManagerEx
 import com.intellij.diff.requests.SimpleDiffRequest
 import com.intellij.ide.BrowserUtil
 import com.intellij.openapi.application.ApplicationManager
-import com.intellij.openapi.fileEditor.FileDocumentManager
 import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.project.guessProjectDir
@@ -473,16 +472,14 @@ class CodeTestChatController(
                     )
                     session.openedDiffFile = FileEditorManager.getInstance(context.project).selectedEditor?.file
                     ApplicationManager.getApplication().runReadAction {
-                        generatedFileContent = getFileContentAtTestFilePath(
-                            session.projectRoot,
-                            session.testFileRelativePathToProjectRoot
-                        )
-                        val selectedFile = FileEditorManager.getInstance(context.project).selectedEditor?.file
-                        selectedFileContent = selectedFile?.let {
-                            FileDocumentManager.getInstance().getDocument(it)?.text
-                        }.orEmpty()
+                        generatedFileContent = getGeneratedFileContent(session)
                     }
 
+                    selectedFileContent = getFileContentAtTestFilePath(
+                        session.projectRoot,
+                        session.testFileRelativePathToProjectRoot,
+                    )
+
                     // Line difference calculation: linesOfCodeGenerated = number of lines in generated test file - number of lines in original test file
                     numberOfLinesGenerated = generatedFileContent.lines().size
                     numberOfLinesSelected = selectedFileContent.lines().size
@@ -1070,6 +1067,12 @@ class CodeTestChatController(
         }
     }
 
+    // Return generated test file content
+    private fun getGeneratedFileContent(session: Session): String {
+        val generateFileContent = session.generatedTestDiffs[session.testFileRelativePathToProjectRoot].toString()
+        return generateFileContent
+    }
+
     /*
      If shortAnswer has buildCommand, use it, if it doesn't hardcode it according to the user type(internal or not)
     private fun getBuildCommand(tabId: String): String {