Merge pull request #303 from JetBrains-Research/ebraun/refactoring/te…

…st-class-builder-helper

[#301] TestClassBuilderHelper refactoring

CONTRIBUTING.md

@@ -1,5 +1,15 @@
 # TestSpark
 
+## Table of contents
+- [Description](#description)
+- [Build project](#build-project)
+- [Run IDE for UI tests](#run-IDE-for-ui-tests)
+- [Plugin Configuration](#plugin-configuration-file)
+- [Language Support Documentation](#language-support-documentation)
+- [Classes](#classes)
+- [Tests](#tests)
+
+
 ## Description
 
 In this document you can find the overall structure of TestSpark plugin. The classes are listed and their purpose is described. This section is intended for developers and contributors to TestSpark plugin.
@@ -23,6 +33,139 @@ to include test generation using Grazie in the runIdeForUiTests process, you nee
 
 `<TOKEN>` is generated by Space, which has access to Automatically generating unit tests maven packages.
 
+---
+
+## Language Support Documentation
+
+The TestSpark plugin supports automatic test generation for various programming languages (currently Java and Kotlin)
+and aims to support even more programming languages in the future.
+
+This document provides an overview of the existing implementation of Kotlin and Java support and guidelines for adding
+more programming languages.
+
+>How can I add support for a new programming language? 
+In brief, you need to extend all the necessary interfaces with implementations specific to the new language.
+Below, you will find a detailed guide divided into six key components of the entire pipeline with the most 
+important interfaces addressing this goal. 
+
+
+## Key Components
+
+### 1. PSI Parsers
+
+The first step is to enable the collection of the appropriate information for the code under test. This part is
+responsible for working with the PSI (Program Structure Interface) generated by IntelliJ IDEA. It helps parse the part
+where the cursor is located, provides a choice of the code elements that are available for testing at cursor's position.
+Then find all the needed dependencies to make the prompt complete with all the necessary knowledge about the code under
+test.
+
+This part is the most granular but complex at the same time.
+
+The main reason for this is to include dependencies only for the languages we need. This avoids errors if the user does
+not have some languages that our plugin supports. _For example, if we work with a Python project, we don't want to depend
+on Kotlin because it will cause an error if Kotlin isn't present._
+
+Additionally, we want to incrementally add dependencies on other languages for faster startup.
+_For example, we do not want to fetch the dependency on Java when we work with TypeScript._
+Other benefits include better organization, easier maintenance, and clearer separation of
+concerns. As a side-bonus, the addition of new languages will be easier.
+
+**Module Dependencies:**
+
+- **langwrappers**: This is a foundational module for language extensions.
+- **<Language>**: Depends on the `langwrappers` module to implement the `<Language>`-specific `PsiHelper`
+  and `PsiHelperProvider`.
+- **src/**: Depends on `langwrappers` because we want to use `PsiHelper` and other interfaces regardless of the current
+  language. Depends on `<Language>`, to make `plugin.xml` aware of the implementations of the Extension Point.
+
+**Plugin Dependencies:**
+
+- The main `plugin.xml` file declares the `psiHelperProvider` extension point using
+  the `com.intellij.lang.LanguageExtensionPoint` class.
+- The language-specific modules extend this extension point to register their implementations.
+- When the project is opened, we load the EPs needed to work with the current project. Then, using
+  the `PsiHelperProvider` interface, we can get the appropriate `<language>PsiHelper` class per file.
+
+**Implementation Details:**
+
+- **Common Module (`langwrappers`)**:
+    - Contains the `PsiHelper` interface, which provides the necessary methods to interact with `psiFile`.
+    - The `PsiHelperProvider` class includes a companion object to fetch the appropriate `PsiHelper` implementation
+      based on the file's language.
+
+- **<Language> Module**:
+    - Implements the `<Language>PsiHelper` and `<Language>PsiHelperProvider` classes, which provide <Language>-specific
+      logic.
+    - Declares the extension point in `testspark-<Language>.xml`.
+
+To add new languages, create a separate module for this language and register its implementation as an extension of
+the `psiHelperProvider` EP. Then follow the template provided above.
+
+### 2. Prompt Generation
+
+When we know how to parse the code, we need to construct the prompt.
+
+For each language, adjust the prompt that goes to the LLM. Ensure that the language, framework platform, and mocking
+framework are defined correctly in:
+
+```kotlin
+data class PromptConfiguration(
+    val desiredLanguage: String,
+    val desiredTestingPlatform: String,
+    val desiredMockingFramework: String,
+)
+```
+
+Additionally, check that all the dependencies (collected by `PsiHelper` for the current strategy) are passed
+properly. `PromptGenerator` and `PromptBuilder` are responsible for this job.
+
+### 3. Parsing LLM Response
+
+When the LLM response to our prompt is received, we have to parse it.
+
+We want to retrieve test case, all the test functions and additional information like imports or supporting functions
+from the response.
+
+The current structure of this part is located in:
+
+- `kotlin/org/jetbrains/research/testspark/core/test`
+- `kotlin/org/jetbrains/research/testspark/tools`
+
+It can be more easily understood with the following diagram:
+![](https://private-user-images.githubusercontent.com/70476032/349256986-dc7e1ff9-a9a5-4bd2-a51f-ecbfabeb6cba.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MjIzNTEyOTAsIm5iZiI6MTcyMjM1MDk5MCwicGF0aCI6Ii83MDQ3NjAzMi8zNDkyNTY5ODYtZGM3ZTFmZjktYTlhNS00YmQyLWE1MWYtZWNiZmFiZWI2Y2JhLnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDA3MzAlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwNzMwVDE0NDk1MFomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPWJjMDg3MWM2ZDA4MDJlZGUwNzliMzNkNzA3YWI4YTcwM2RmYTFjMmE1MGM4MjM5NjJiOGI2ZjgxNTE2OTU2YjQmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.8OfRa1wJhDfFq3QT6h5yIjBh1VqB9UrrQfZGp0_SLDo)
+
+- `TestsAssembler`: Assembler class for generating and organizing test cases from the LLM response.
+- `TestSuiteParser`: Extracts test cases from raw text and generates a test suite.
+- `TestBodyPrinter`: Generates the body of a test function as a string.
+
+### 4. Compilation
+
+Before showing the code to the user, it should be checked for compilation.
+
+- `TestCompiler`: Compiles a list of test cases and returns the compilation result.
+
+Here one should specify the appropriate compilation strategy for each language. With all the dependencies and build paths.
+
+### 5. UI Representation
+
+Once the code generated by the LLM is checked for the compilation, it should be presented in the UI.
+
+- `TestCaseDisplayService`: Service responsible for the representation of all the UI components.
+- `TestSuiteView`: Interface specific for working with buttons.
+- `TestClassCodeAnalyzer`: Interface for retrieving information from test class code.
+- `TestClassCodeGenerator`: Interface for generating and formatting test class code.
+
+### 6. Running and saving tests
+
+We should be able to run all the tests in the UI and then save them to the desired folder.
+
+- `TestPersistentStorage`: Interface representing a contract for saving generated tests to a specified file system location.
+
+For Kotlin and Java, the `TestProcessor` implementation also allows saving the JaCoCo report to see the code coverage of 
+the test that will be saved.
+
+---
+
 ## Plugin Configuration File
 
 The plugin configuration file is `plugin.xml` which can be found in `src/main/resources/META-INF` directory. All declarations (such as actions, services, listeners) are present in this file.

core/src/main/kotlin/org/jetbrains/research/testspark/core/test/TestsPersistentStorage.kt

@@ -4,6 +4,7 @@ package org.jetbrains.research.testspark.core.test
  * The TestPersistentStorage interface represents a contract for saving generated tests to a specified file system location.
  */
 interface TestsPersistentStorage {
+
     /**
      * Save the generated tests to a specified directory.
      *

core/src/main/kotlin/org/jetbrains/research/testspark/core/test/kotlin/KotlinTestCompiler.kt

@@ -12,6 +12,7 @@ class KotlinTestCompiler(libPaths: List<String>, junitLibPaths: List<String>) :
     override fun compileCode(path: String, projectBuildPath: String): Pair<Boolean, String> {
         log.info { "[KotlinTestCompiler] Compiling ${path.substringAfterLast('/')}" }
 
+        // TODO find the kotlinc if it is not in PATH
         val classPaths = "\"${getClassPaths(projectBuildPath)}\""
         // Compile file
         val errorMsg = CommandLineRunner.run(
@@ -23,7 +24,12 @@ class KotlinTestCompiler(libPaths: List<String>, junitLibPaths: List<String>) :
             ),
         )
 
-        log.info { "Error message: '$errorMsg'" }
+        if (errorMsg.isNotEmpty()) {
+            log.info { "Error message: '$errorMsg'" }
+            if (errorMsg.contains("kotlinc: command not found'")) {
+                throw RuntimeException(errorMsg)
+            }
+        }
 
         // No need to save the .class file for kotlin, so checking the error message is enough
         return Pair(errorMsg.isBlank(), errorMsg)

java/src/main/kotlin/org/jetbrains/research/testspark/java/JavaPsiHelper.kt

@@ -67,7 +67,7 @@ class JavaPsiHelper(private val psiFile: PsiFile) : PsiHelper {
         return null
     }
 
-    override fun getSurroundingLine(caretOffset: Int): Int? {
+    override fun getSurroundingLineNumber(caretOffset: Int): Int? {
         val doc = PsiDocumentManager.getInstance(psiFile.project).getDocument(psiFile) ?: return null
 
         val selectedLine = doc.getLineNumber(caretOffset)
@@ -158,7 +158,7 @@ class JavaPsiHelper(private val psiFile: PsiFile) : PsiHelper {
 
         val javaPsiClassWrapped = getSurroundingClass(caret.offset) as JavaPsiClassWrapper?
         val javaPsiMethodWrapped = getSurroundingMethod(caret.offset) as JavaPsiMethodWrapper?
-        val line: Int? = getSurroundingLine(caret.offset)
+        val line: Int? = getSurroundingLineNumber(caret.offset)
 
         javaPsiClassWrapped?.let { result.add(CodeType.CLASS to getClassHTMLDisplayName(it)) }
         javaPsiMethodWrapped?.let { result.add(CodeType.METHOD to getMethodHTMLDisplayName(it)) }

kotlin/src/main/kotlin/org/jetbrains/research/testspark/kotlin/KotlinPsiHelper.kt

@@ -7,22 +7,13 @@ import com.intellij.openapi.editor.Caret
 import com.intellij.openapi.module.ModuleUtilCore
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.TextRange
-import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiFile
 import com.intellij.psi.util.parentOfType
-import org.jetbrains.kotlin.asJava.toLightClass
-import org.jetbrains.kotlin.descriptors.ClassDescriptor
-import org.jetbrains.kotlin.idea.base.psi.kotlinFqName
-import org.jetbrains.kotlin.idea.caches.resolve.analyze
-import org.jetbrains.kotlin.psi.KtClass
 import org.jetbrains.kotlin.psi.KtClassOrObject
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtFunction
-import org.jetbrains.kotlin.psi.KtTypeReference
-import org.jetbrains.kotlin.resolve.BindingContext
-import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
-import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
+import org.jetbrains.kotlin.psi.KtPsiUtil
 import org.jetbrains.research.testspark.core.test.SupportedLanguage
 import org.jetbrains.research.testspark.core.test.data.CodeType
 import org.jetbrains.research.testspark.langwrappers.CodeTypeDisplayName
@@ -70,7 +61,7 @@ class KotlinPsiHelper(private val psiFile: PsiFile) : PsiHelper {
         return null
     }
 
-    override fun getSurroundingLine(caretOffset: Int): Int? {
+    override fun getSurroundingLineNumber(caretOffset: Int): Int? {
         val doc = PsiDocumentManager.getInstance(psiFile.project).getDocument(psiFile) ?: return null
 
         val selectedLine = doc.getLineNumber(caretOffset)
@@ -121,19 +112,13 @@ class KotlinPsiHelper(private val psiFile: PsiFile) : PsiHelper {
 
         repeat(maxInputParamsDepth) {
             val tempListOfClasses = mutableSetOf<KotlinPsiClassWrapper>()
-
             currentLevelClasses.forEach { classIt ->
                 classIt.methods.forEach { methodIt ->
                     (methodIt as KotlinPsiMethodWrapper).parameterList?.parameters?.forEach { paramIt ->
-                        val typeRef = paramIt.typeReference
-                        if (typeRef != null) {
-                            resolveClassInType(typeRef)?.let { psiClass ->
-                                if (psiClass.kotlinFqName != null) {
-                                    KotlinPsiClassWrapper(psiClass as KtClass).let {
-                                        if (!it.qualifiedName.startsWith("kotlin.")) {
-                                            interestingPsiClasses.add(it)
-                                        }
-                                    }
+                        KtPsiUtil.getClassIfParameterIsProperty(paramIt)?.let { typeIt ->
+                            KotlinPsiClassWrapper(typeIt).let {
+                                if (!it.qualifiedName.startsWith("kotlin.")) {
+                                    interestingPsiClasses.add(it)
                                 }
                             }
                         }
@@ -165,7 +150,7 @@ class KotlinPsiHelper(private val psiFile: PsiFile) : PsiHelper {
 
         val ktClass = getSurroundingClass(caret.offset)
         val ktFunction = getSurroundingMethod(caret.offset)
-        val line: Int? = getSurroundingLine(caret.offset)?.plus(1)
+        val line: Int? = getSurroundingLineNumber(caret.offset)?.plus(1)
 
         ktClass?.let { result.add(CodeType.CLASS to getClassHTMLDisplayName(it)) }
         ktFunction?.let { result.add(CodeType.METHOD to getMethodHTMLDisplayName(it)) }
@@ -202,11 +187,4 @@ class KotlinPsiHelper(private val psiFile: PsiFile) : PsiHelper {
             else -> "<html><b><font color='orange'>method</font> ${psiMethod.name}</b></html>"
         }
     }
-
-    private fun resolveClassInType(typeReference: KtTypeReference): PsiClass? {
-        val context = typeReference.analyze(BodyResolveMode.PARTIAL)
-        val type = context[BindingContext.TYPE, typeReference] ?: return null
-        val classDescriptor = type.constructor.declarationDescriptor as? ClassDescriptor ?: return null
-        return (DescriptorToSourceUtils.getSourceFromDescriptor(classDescriptor) as? KtClass)?.toLightClass()
-    }
 }

langwrappers/src/main/kotlin/org/jetbrains/research/testspark/langwrappers/PsiComponents.kt

@@ -114,7 +114,7 @@ interface PsiHelper {
      * @param caretOffset The caret offset within the PSI file.
      * @return The line number of the selected line, otherwise null.
      */
-    fun getSurroundingLine(caretOffset: Int): Int?
+    fun getSurroundingLineNumber(caretOffset: Int): Int?
 
     /**
      * Retrieves a set of interesting PsiClasses based on a given project,

src/main/kotlin/org/jetbrains/research/testspark/appstarter/TestSparkStarter.kt

@@ -174,7 +174,7 @@ class TestSparkStarter : ApplicationStarter {
                         // Start test generation
                         val indicator = HeadlessProgressIndicator()
                         val errorMonitor = DefaultErrorMonitor()
-                        val testCompiler = TestCompilerFactory.createTestCompiler(
+                        val testCompiler = TestCompilerFactory.create(
                             project,
                             settingsState.junitVersion,
                             psiHelper.language,

src/main/kotlin/org/jetbrains/research/testspark/display/TestCasePanelFactory.kt

@@ -32,13 +32,12 @@ import org.jetbrains.research.testspark.data.llm.JsonEncoding
 import org.jetbrains.research.testspark.display.custom.IJProgressIndicator
 import org.jetbrains.research.testspark.helpers.LLMHelper
 import org.jetbrains.research.testspark.helpers.ReportHelper
-import org.jetbrains.research.testspark.helpers.java.JavaClassBuilderHelper
-import org.jetbrains.research.testspark.helpers.kotlin.KotlinClassBuilderHelper
 import org.jetbrains.research.testspark.services.LLMSettingsService
 import org.jetbrains.research.testspark.services.TestsExecutionResultService
 import org.jetbrains.research.testspark.services.java.JavaTestCaseDisplayService
 import org.jetbrains.research.testspark.services.kotlin.KotlinTestCaseDisplayService
 import org.jetbrains.research.testspark.settings.llm.LLMSettingsState
+import org.jetbrains.research.testspark.tools.TestClassCodeAnalyzerFactory
 import org.jetbrains.research.testspark.tools.TestCompilerFactory
 import org.jetbrains.research.testspark.tools.TestProcessor
 import org.jetbrains.research.testspark.tools.ToolUtils
@@ -469,17 +468,7 @@ class TestCasePanelFactory(
         WriteCommandAction.runWriteCommandAction(project) {
             uiContext.errorMonitor.clear()
             val code = testSuitePresenter.toString(testSuite)
-            testCase.testName = when (language) {
-                SupportedLanguage.Kotlin -> KotlinClassBuilderHelper.extractFirstTestMethodName(
-                    testCase.testName,
-                    code,
-                )
-
-                SupportedLanguage.Java -> JavaClassBuilderHelper.extractFirstTestMethodName(
-                    testCase.testName,
-                    code,
-                )
-            }
+            testCase.testName = TestClassCodeAnalyzerFactory.create(language).extractFirstTestMethodName(testCase.testName, code)
             testCase.testCode = code
 
             // update numbers
@@ -537,15 +526,9 @@ class TestCasePanelFactory(
     private fun runTest(indicator: CustomProgressIndicator) {
         indicator.setText("Executing ${testCase.testName}")
 
-        val fileName = when (language) {
-            SupportedLanguage.Kotlin ->
-                "${KotlinClassBuilderHelper.getClassFromTestCaseCode(testCase.testCode)}.kt"
-
-            SupportedLanguage.Java ->
-                "${JavaClassBuilderHelper.getClassFromTestCaseCode(testCase.testCode)}.java"
-        }
+        val fileName = TestClassCodeAnalyzerFactory.create(language).getFileNameFromTestCaseCode(testCase.testName)
 
-        val testCompiler = TestCompilerFactory.createTestCompiler(
+        val testCompiler = TestCompilerFactory.create(
             project,
             llmSettingsState.junitVersion,
             language,
@@ -708,17 +691,7 @@ class TestCasePanelFactory(
      * Updates the current test case with the specified test name and test code.
      */
     private fun updateTestCaseInformation() {
-        testCase.testName = when (language) {
-            SupportedLanguage.Kotlin -> KotlinClassBuilderHelper.extractFirstTestMethodName(
-                testCase.testName,
-                languageTextField.document.text,
-            )
-
-            SupportedLanguage.Java -> JavaClassBuilderHelper.extractFirstTestMethodName(
-                testCase.testName,
-                languageTextField.document.text,
-            )
-        }
+        testCase.testName = TestClassCodeAnalyzerFactory.create(language).extractFirstTestMethodName(testCase.testName, languageTextField.document.text)
         testCase.testCode = languageTextField.document.text
     }
 }