Cleanup CodeWhispererCodeFileScanTest setup (#4863)

25 lateinit vars numbered 1/2/3/4/5 is not readable

Author: rli

plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeFileScanTest.kt

@@ -5,21 +5,15 @@ package software.aws.toolkits.jetbrains.services.codewhisperer.codescan
 
 import com.intellij.openapi.fileEditor.FileEditor
 import com.intellij.openapi.fileEditor.FileEditorManager
-import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiFile
-import com.intellij.testFramework.replaceService
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.runTest
 import org.apache.commons.codec.digest.DigestUtils
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.Before
 import org.junit.Test
 import org.mockito.ArgumentMatchers.anyString
-import org.mockito.Mockito
 import org.mockito.Mockito.mock
-import org.mockito.Mockito.times
-import org.mockito.Mockito.`when`
-import org.mockito.internal.verification.Times
 import org.mockito.kotlin.any
 import org.mockito.kotlin.argumentCaptor
 import org.mockito.kotlin.doNothing
@@ -28,7 +22,9 @@ import org.mockito.kotlin.inOrder
 import org.mockito.kotlin.isNull
 import org.mockito.kotlin.spy
 import org.mockito.kotlin.stub
+import org.mockito.kotlin.times
 import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
 import software.amazon.awssdk.awscore.exception.AwsErrorDetails
 import software.amazon.awssdk.services.codewhisperer.model.CodeWhispererException
 import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlRequest
@@ -42,7 +38,6 @@ import software.aws.toolkits.jetbrains.utils.isInstanceOf
 import software.aws.toolkits.jetbrains.utils.isInstanceOfSatisfying
 import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule
 import software.aws.toolkits.telemetry.CodewhispererLanguage
-import java.io.File
 import java.io.FileInputStream
 import java.lang.management.ManagementFactory
 import java.util.Base64
@@ -52,55 +47,19 @@ import kotlin.io.path.relativeTo
 import kotlin.test.assertNotNull
 
 class CodeWhispererCodeFileScanTest : CodeWhispererCodeScanTestBase(PythonCodeInsightTestFixtureRule()) {
-    private lateinit var psifile: PsiFile
-    private lateinit var psifile2: PsiFile
-    private lateinit var psifile3: PsiFile
-    private lateinit var psifile4: PsiFile
-    private lateinit var psifilePerformanceTest: PsiFile
-    private lateinit var psifilePerformanceTest2: PsiFile
-    private lateinit var file: File
-    private lateinit var file2: File
-    private lateinit var file3: File
-    private lateinit var file4: File
-    private lateinit var performanceTestfileWithPayload200KB: File
-    private lateinit var performanceTestfileWithPayload150KB: File
-    private lateinit var virtualFile3: VirtualFile
-    private lateinit var virtualFile4: VirtualFile
-    private lateinit var sessionConfigSpy: CodeScanSessionConfig
-    private lateinit var sessionConfigSpy2: CodeScanSessionConfig
-    private lateinit var sessionConfigSpy3: CodeScanSessionConfig
-    private lateinit var sessionConfigSpy4: CodeScanSessionConfig
+    private val codeScanName = UUID.randomUUID().toString()
     private val payloadContext = PayloadContext(CodewhispererLanguage.Python, 1, 1, 10, listOf(), 600, 200)
-    private lateinit var codeScanSessionContext: CodeScanSessionContext
-    private lateinit var codeScanSessionContext2: CodeScanSessionContext
-    private lateinit var codeScanSessionContext3: CodeScanSessionContext
+
+    private lateinit var pyPsiFile: PsiFile
+    private lateinit var ktPsiFile: PsiFile
+    private lateinit var pySession: CodeScanSessionConfig
     private lateinit var codeScanSessionSpy: CodeWhispererCodeScanSession
-    private lateinit var codeScanSessionSpy2: CodeWhispererCodeScanSession
-    private lateinit var codeScanSessionSpy3: CodeWhispererCodeScanSession
-    private val codeScanName = UUID.randomUUID().toString()
 
     @Before
     override fun setup() {
         super.setup()
 
-        psifile2 = projectRule.fixture.addFileToProject(
-            "/subtract.java",
-            """public class MathOperations {
-                public static int subtract(int a, int b) {
-                    return a - b; 
-                    }
-                public static void main(String[] args) {    
-                    int num1 = 10;
-                    int num2 = 5;
-                    int result = subtract(num1, num2);
-                    System.out.println(result);
-                    }
-                }     
-            """.trimMargin()
-        )
-        file2 = psifile2.virtualFile.toNioPath().toFile()
-
-        psifile = projectRule.fixture.addFileToProject(
+        pyPsiFile = projectRule.fixture.addFileToProject(
             "/test.py",
             """import numpy as np
                import from module1 import helper
@@ -110,9 +69,8 @@ class CodeWhispererCodeFileScanTest : CodeWhispererCodeScanTestBase(PythonCodeIn
                   
             """.trimMargin()
         )
-        file = psifile.virtualFile.toNioPath().toFile()
 
-        psifile3 = projectRule.fixture.addFileToProject(
+        ktPsiFile = projectRule.fixture.addFileToProject(
             "/test.kt",
             // write simple addition function in kotlin
             """
@@ -124,96 +82,68 @@ class CodeWhispererCodeFileScanTest : CodeWhispererCodeScanTestBase(PythonCodeIn
                 }
             """.trimMargin()
         )
-        virtualFile3 = psifile3.virtualFile
-        file3 = virtualFile3.toNioPath().toFile()
 
-        psifile4 = projectRule.fixture.addFileToProject(
-            "../test.java",
-            """
-                public class Addition {
-                    public static void main(String[] args) {
-                        int a = 1;
-                        int b = 2;
-                        int c = a + b;
-                        System.out.println(c);
+        projectRule.fixture.addFileToProject(
+            "/subtract.java",
+            """public class MathOperations {
+                public static int subtract(int a, int b) {
+                    return a - b; 
                     }
-                }
-                """
-        )
-        virtualFile4 = psifile4.virtualFile
-        file4 = virtualFile4.toNioPath().toFile()
-
-        // Create a 200KB file
-        val content = "a".repeat(200 * 1024)
-        psifilePerformanceTest = projectRule.fixture.addFileToProject("test.txt", content)
-        performanceTestfileWithPayload200KB = psifilePerformanceTest.virtualFile.toNioPath().toFile()
-
-        sessionConfigSpy3 = spy(
-            CodeScanSessionConfig.create(
-                psifilePerformanceTest.virtualFile,
-                project,
-                CodeWhispererConstants.CodeAnalysisScope.FILE
-            )
-        )
-        setupResponse(psifilePerformanceTest.virtualFile.toNioPath().relativeTo(sessionConfigSpy3.projectRoot.toNioPath()))
-
-        // Create a 150KB file
-        val codeContentForPayload = "a".repeat(150 * 1024)
-        psifilePerformanceTest2 = projectRule.fixture.addFileToProject("test.txt", codeContentForPayload)
-        performanceTestfileWithPayload150KB = psifilePerformanceTest2.virtualFile.toNioPath().toFile()
-
-        sessionConfigSpy4 = spy(
-            CodeScanSessionConfig.create(
-                psifilePerformanceTest2.virtualFile,
-                project,
-                CodeWhispererConstants.CodeAnalysisScope.FILE
-            )
-        )
-        setupResponse(psifilePerformanceTest2.virtualFile.toNioPath().relativeTo(sessionConfigSpy4.projectRoot.toNioPath()))
-        sessionConfigSpy = spy(
-            CodeScanSessionConfig.create(
-                psifile.virtualFile,
-                project,
-                CodeWhispererConstants.CodeAnalysisScope.FILE
-            )
+                public static void main(String[] args) {    
+                    int num1 = 10;
+                    int num2 = 5;
+                    int result = subtract(num1, num2);
+                    System.out.println(result);
+                    }
+                }     
+            """.trimMargin()
         )
 
-        sessionConfigSpy2 = spy(
+        pySession = spy(
             CodeScanSessionConfig.create(
-                virtualFile4,
+                pyPsiFile.virtualFile,
                 project,
                 CodeWhispererConstants.CodeAnalysisScope.FILE
             )
         )
+        setupResponse(pyPsiFile.virtualFile.toNioPath().relativeTo(pySession.projectRoot.toNioPath()))
 
-        setupResponse(psifile.virtualFile.toNioPath().relativeTo(sessionConfigSpy.projectRoot.toNioPath()))
-
-        sessionConfigSpy.stub {
-            onGeneric { sessionConfigSpy.createPayload() }.thenReturn(Payload(payloadContext, file))
+        pySession.stub {
+            onGeneric { pySession.createPayload() }.thenReturn(Payload(payloadContext, pyPsiFile.virtualFile.toNioPath().toFile()))
         }
 
         // Mock CodeWhispererClient needs to be setup before initializing CodeWhispererCodeScanSession
-        codeScanSessionContext = CodeScanSessionContext(project, sessionConfigSpy, CodeWhispererConstants.CodeAnalysisScope.FILE)
-        codeScanSessionSpy = spy(CodeWhispererCodeScanSession(codeScanSessionContext))
-        doNothing().`when`(codeScanSessionSpy).uploadArtifactToS3(any(), any(), any(), any(), isNull(), any())
-
-        codeScanSessionContext2 = CodeScanSessionContext(project, sessionConfigSpy3, CodeWhispererConstants.CodeAnalysisScope.FILE)
-        codeScanSessionSpy2 = spy(CodeWhispererCodeScanSession(codeScanSessionContext2))
-        doNothing().`when`(codeScanSessionSpy2).uploadArtifactToS3(any(), any(), any(), any(), isNull(), any())
+        val pySessionContext = CodeScanSessionContext(project, pySession, CodeWhispererConstants.CodeAnalysisScope.FILE)
+        codeScanSessionSpy = spy(CodeWhispererCodeScanSession(pySessionContext))
+        doNothing().whenever(codeScanSessionSpy).uploadArtifactToS3(any(), any(), any(), any(), isNull(), any())
 
-        codeScanSessionContext3 = CodeScanSessionContext(project, sessionConfigSpy4, CodeWhispererConstants.CodeAnalysisScope.FILE)
-        codeScanSessionSpy3 = spy(CodeWhispererCodeScanSession(codeScanSessionContext3))
-        doNothing().`when`(codeScanSessionSpy3).uploadArtifactToS3(any(), any(), any(), any(), isNull(), any())
         mockClient.stub {
-            onGeneric { createUploadUrl(any()) }.thenReturn(fakeCreateUploadUrlResponse)
-            onGeneric { createCodeScan(any(), any()) }.thenReturn(fakeCreateCodeScanResponse)
-            onGeneric { getCodeScan(any(), any()) }.thenReturn(fakeGetCodeScanResponse)
-            onGeneric { listCodeScanFindings(any(), any()) }.thenReturn(fakeListCodeScanFindingsResponse)
+            // setupResponse dynamically modifies these fake responses so this is very hard to follow and makes me question if we even need this
+            onGeneric { createUploadUrl(any()) }.thenAnswer { fakeCreateUploadUrlResponse }
+            onGeneric { createCodeScan(any(), any()) }.thenAnswer { fakeCreateCodeScanResponse }
+            onGeneric { getCodeScan(any(), any()) }.thenAnswer { fakeGetCodeScanResponse }
+            onGeneric { listCodeScanFindings(any(), any()) }.thenAnswer { fakeListCodeScanFindingsResponse }
         }
     }
 
     @Test
-    fun `test run() - measure CPU and memory usage`() {
+    fun `test run() - measure CPU and memory usage with payload of 200KB`() {
+        // Create a 200KB file
+        val content = "a".repeat(200 * 1024)
+        val psiFile = projectRule.fixture.addFileToProject("test.txt", content)
+
+        val sessionConfig = spy(
+            CodeScanSessionConfig.create(
+                psiFile.virtualFile,
+                project,
+                CodeWhispererConstants.CodeAnalysisScope.FILE
+            )
+        )
+        setupResponse(psiFile.virtualFile.toNioPath().relativeTo(sessionConfig.projectRoot.toNioPath()))
+        val sessionContext = CodeScanSessionContext(project, sessionConfig, CodeWhispererConstants.CodeAnalysisScope.FILE)
+        val session = spy(CodeWhispererCodeScanSession(sessionContext))
+        doNothing().whenever(session).uploadArtifactToS3(any(), any(), any(), any(), isNull(), any())
+
         // Set up CPU and Memory monitoring
         val runtime = Runtime.getRuntime()
         val bean = ManagementFactory.getThreadMXBean()
@@ -223,7 +153,7 @@ class CodeWhispererCodeFileScanTest : CodeWhispererCodeScanTestBase(PythonCodeIn
 
         // Run the code scan
         runBlocking {
-            codeScanSessionSpy2.run()
+            session.run()
         }
 
         // Calculate CPU and memory usage
@@ -248,6 +178,22 @@ class CodeWhispererCodeFileScanTest : CodeWhispererCodeScanTestBase(PythonCodeIn
 
     @Test
     fun `test run() - measure CPU and memory usage with payload of 150KB`() {
+        // Create a 150KB file
+        val codeContentForPayload = "a".repeat(150 * 1024)
+        val psiFile = projectRule.fixture.addFileToProject("test.txt", codeContentForPayload)
+
+        val sessionConfig = spy(
+            CodeScanSessionConfig.create(
+                psiFile.virtualFile,
+                project,
+                CodeWhispererConstants.CodeAnalysisScope.FILE
+            )
+        )
+        setupResponse(psiFile.virtualFile.toNioPath().relativeTo(sessionConfig.projectRoot.toNioPath()))
+        val sessionContext = CodeScanSessionContext(project, sessionConfig, CodeWhispererConstants.CodeAnalysisScope.FILE)
+        val session = spy(CodeWhispererCodeScanSession(sessionContext))
+        doNothing().whenever(session).uploadArtifactToS3(any(), any(), any(), any(), isNull(), any())
+
         // Set up CPU and Memory monitoring
         val runtime = Runtime.getRuntime()
         val bean = ManagementFactory.getThreadMXBean()
@@ -257,7 +203,7 @@ class CodeWhispererCodeFileScanTest : CodeWhispererCodeScanTestBase(PythonCodeIn
 
         // Run the code scan
         runBlocking {
-            codeScanSessionSpy3.run()
+            session.run()
         }
 
         // Calculate CPU and memory usage
@@ -282,6 +228,7 @@ class CodeWhispererCodeFileScanTest : CodeWhispererCodeScanTestBase(PythonCodeIn
 
     @Test
     fun `test createUploadUrlAndUpload()`() {
+        val file = pyPsiFile.virtualFile.toNioPath().toFile()
         val fileMd5: String = Base64.getEncoder().encodeToString(DigestUtils.md5(FileInputStream(file)))
         codeScanSessionSpy.stub {
             onGeneric { codeScanSessionSpy.createUploadUrl(any(), any(), any()) }
@@ -327,7 +274,9 @@ class CodeWhispererCodeFileScanTest : CodeWhispererCodeScanTestBase(PythonCodeIn
 
     @Test
     fun `test run() - happypath`() = runTest {
-        assertNotNull(sessionConfigSpy)
+        val file = pyPsiFile.virtualFile.toNioPath().toFile()
+
+        assertNotNull(pySession)
         val codeScanResponse = codeScanSessionSpy.run()
         assertThat(codeScanResponse).isInstanceOfSatisfying<CodeScanResponse.Success> {
             assertThat(it.issues).hasSize(2)
@@ -336,14 +285,38 @@ class CodeWhispererCodeFileScanTest : CodeWhispererCodeScanTestBase(PythonCodeIn
         }
 
         val inOrder = inOrder(codeScanSessionSpy)
-        inOrder.verify(codeScanSessionSpy, Times(1)).createUploadUrlAndUpload(eq(file), eq("SourceCode"), anyString())
-        inOrder.verify(codeScanSessionSpy, Times(1)).createCodeScan(eq(CodewhispererLanguage.Python.toString()), anyString())
-        inOrder.verify(codeScanSessionSpy, Times(1)).getCodeScan(any())
-        inOrder.verify(codeScanSessionSpy, Times(1)).listCodeScanFindings(eq("jobId"), eq(null))
+        inOrder.verify(codeScanSessionSpy, times(1)).createUploadUrlAndUpload(eq(file), eq("SourceCode"), anyString())
+        inOrder.verify(codeScanSessionSpy, times(1)).createCodeScan(eq(CodewhispererLanguage.Python.toString()), anyString())
+        inOrder.verify(codeScanSessionSpy, times(1)).getCodeScan(any())
+        inOrder.verify(codeScanSessionSpy, times(1)).listCodeScanFindings(eq("jobId"), eq(null))
     }
 
     @Test
     fun `test createPayload for files outside Project Root`() {
+        val externalFile = projectRule.fixture.addFileToProject(
+            "../test.java",
+            """
+            public class Addition {
+                public static void main(String[] args) {
+                    int a = 1;
+                    int b = 2;
+                    int c = a + b;
+                    System.out.println(c);
+                }
+            }
+            """
+        )
+
+        val sessionConfigSpy2 = spy(
+            CodeScanSessionConfig.create(
+                externalFile.virtualFile,
+                project,
+                CodeWhispererConstants.CodeAnalysisScope.FILE
+            )
+        )
+
+        setupResponse(pyPsiFile.virtualFile.toNioPath().relativeTo(pySession.projectRoot.toNioPath()))
+
         val payload = sessionConfigSpy2.createPayload()
         assertNotNull(payload)
         val payloadZipFile = ZipFile(payload.srcZip)
@@ -354,16 +327,13 @@ class CodeWhispererCodeFileScanTest : CodeWhispererCodeScanTestBase(PythonCodeIn
 
     @Test
     fun `unsupported languages file scan fail`() = runTest {
-        scanManagerSpy = Mockito.spy(CodeWhispererCodeScanManager.getInstance(projectRule.project))
-        projectRule.project.replaceService(CodeWhispererCodeScanManager::class.java, scanManagerSpy, disposableRule.disposable)
-
-        val fileEditorManager = mock(FileEditorManager::class.java)
-        val selectedEditor = mock(FileEditor::class.java)
-        val editorList: MutableList<FileEditor> = mutableListOf(selectedEditor)
+        val fileEditorManager = mock<FileEditorManager>()
+        val selectedEditor = mock<FileEditor>()
+        val editorList = mutableListOf(selectedEditor)
 
-        `when`(fileEditorManager.selectedEditorWithRemotes).thenReturn(editorList)
-        `when`(fileEditorManager.selectedEditor).thenReturn(selectedEditor)
-        `when`(selectedEditor.file).thenReturn(virtualFile3)
+        whenever(fileEditorManager.selectedEditorWithRemotes).thenReturn(editorList)
+        whenever(fileEditorManager.selectedEditor).thenReturn(selectedEditor)
+        whenever(selectedEditor.file).thenReturn(ktPsiFile.virtualFile)
 
         scanManagerSpy.runCodeScan(CodeWhispererConstants.CodeAnalysisScope.FILE)
         // verify that function was run but none of the results/error handling methods were called.
@@ -374,7 +344,7 @@ class CodeWhispererCodeFileScanTest : CodeWhispererCodeScanTestBase(PythonCodeIn
 
     @Test
     fun `test run() - file scans limit reached`() = runTest {
-        assertNotNull(sessionConfigSpy)
+        assertNotNull(pySession)
 
         mockClient.stub {
             onGeneric { codeScanSessionSpy.createUploadUrlAndUpload(any(), any(), any()) }.thenThrow(
@@ -441,7 +411,7 @@ class CodeWhispererCodeFileScanTest : CodeWhispererCodeScanTestBase(PythonCodeIn
 
     @Test
     fun `test run() - getCodeScan pending timeout`() = runTest {
-        sessionConfigSpy.stub {
+        pySession.stub {
             onGeneric { overallJobTimeoutInSeconds() }.thenReturn(5)
         }
         mockClient.stub {