Add restore test case button

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

@@ -6,6 +6,9 @@ object TestSparkIcons {
     @JvmField
     val remove = IconLoader.getIcon("/icons/remove.svg", javaClass)
 
+    @JvmField
+    val undoRemove = IconLoader.getIcon("/icons/undo_remove.svg", javaClass)
+
     @JvmField
     val reset = IconLoader.getIcon("/icons/reset.svg", javaClass)
 

src/main/kotlin/org/jetbrains/research/testspark/display/generatedTests/GenerateTestsTabHelper.kt

@@ -2,11 +2,11 @@ package org.jetbrains.research.testspark.display.generatedTests
 
 object GenerateTestsTabHelper {
     /**
-     * A helper method to remove a test case from the cache and from the UI.
+     * A helper method to purge a test case from the cache and from the UI.
      *
      * @param testCaseName the name of the test
      */
-    fun removeTestCase(testCaseName: String, generatedTestsTabData: GeneratedTestsTabData) {
+    fun purgeTestCase(testCaseName: String, generatedTestsTabData: GeneratedTestsTabData) {
         // Update the number of selected test cases if necessary
         if (generatedTestsTabData.testCaseNameToSelectedCheckbox[testCaseName]!!.isSelected) {
             generatedTestsTabData.testsSelected--
@@ -25,6 +25,40 @@ object GenerateTestsTabHelper {
         generatedTestsTabData.testCaseNameToEditorTextField.remove(testCaseName)
     }
 
+    /**
+     * A helper method to remove a test case temporarily.
+     *
+     * @param testCaseName the name of the test to remove
+     */
+    fun removeTestCase(testCaseName: String, generatedTestsTabData: GeneratedTestsTabData) {
+        // Uncheck the selected checkbox
+        generatedTestsTabData.testCaseNameToSelectedCheckbox[testCaseName]!!.isSelected = false
+
+        // Disable the selected checkbox
+        generatedTestsTabData.testCaseNameToSelectedCheckbox[testCaseName]!!.isEnabled = false
+
+        // Update status
+        generatedTestsTabData.testCaseNameToEnabled[testCaseName] = false
+        generatedTestsTabData.testsRemoved++
+    }
+
+    /**
+     * A helper method to restore a removed test case.
+     *
+     * @param testCaseName the name of the test to restore
+     */
+    fun undoRemoveTestCase(testCaseName: String, generatedTestsTabData: GeneratedTestsTabData) {
+        // Check the selected checkbox
+        generatedTestsTabData.testCaseNameToSelectedCheckbox[testCaseName]!!.isSelected = true
+
+        // Enable the selected checkbox
+        generatedTestsTabData.testCaseNameToSelectedCheckbox[testCaseName]!!.isEnabled = true
+
+        // Update status
+        generatedTestsTabData.testCaseNameToEnabled[testCaseName] = true
+        generatedTestsTabData.testsRemoved--
+    }
+
     /**
      * Updates the user interface of the tool window.
      *

src/main/kotlin/org/jetbrains/research/testspark/display/generatedTests/GeneratedTestsTabBuilder.kt

@@ -166,6 +166,7 @@ class GeneratedTestsTabBuilder(
             generatedTestsTabData.allTestCasePanel.add(testCasePanel)
             addSeparator()
 
+            generatedTestsTabData.testCaseNameToEnabled[testCase.testName] = true
             generatedTestsTabData.testCaseNameToPanel[testCase.testName] = testCasePanel
             generatedTestsTabData.testCaseNameToSelectedCheckbox[testCase.testName] = checkbox
             generatedTestsTabData.testCaseNameToEditorTextField[testCase.testName] =
@@ -243,7 +244,7 @@ class GeneratedTestsTabBuilder(
      */
     fun clear() {
         generatedTestsTabData.testCaseNameToPanel.toMap()
-            .forEach { GenerateTestsTabHelper.removeTestCase(it.key, generatedTestsTabData) }
+            .forEach { GenerateTestsTabHelper.purgeTestCase(it.key, generatedTestsTabData) }
         generatedTestsTabData.testCasePanelFactories.clear()
         generatedTestsTabData.topButtonsPanelBuilder.clear(generatedTestsTabData)
 

src/main/kotlin/org/jetbrains/research/testspark/display/generatedTests/GeneratedTestsTabData.kt

@@ -12,7 +12,9 @@ class GeneratedTestsTabData {
     val testCaseNameToPanel: HashMap<String, JPanel> = HashMap()
     val testCaseNameToSelectedCheckbox: HashMap<String, JCheckBox> = HashMap()
     val testCaseNameToEditorTextField: HashMap<String, EditorTextField> = HashMap()
+    val testCaseNameToEnabled: HashMap<String, Boolean> = HashMap()
     var testsSelected: Int = 0
+    var testsRemoved: Int = 0
     val unselectedTestCases: HashMap<Int, TestCase> = HashMap()
     val testCasePanelFactories: ArrayList<TestCasePanelBuilder> = arrayListOf()
     var allTestCasePanel: JPanel = JPanel()

src/main/kotlin/org/jetbrains/research/testspark/display/generatedTests/TestCasePanelBuilder.kt

@@ -118,10 +118,14 @@ class TestCasePanelBuilder(
         ScrollPaneConstants.HORIZONTAL_SCROLLBAR_ALWAYS,
     )
 
-    // Create "Remove" button to remove the test from cache
+    // Create "Remove" button to remove the test temporarily
     private val removeButton =
         IconButtonCreator.getButton(TestSparkIcons.remove, PluginLabelsBundle.get("removeTip"))
 
+    // Create "Undo remove" button to restore a removed test
+    private val undoRemoveButton =
+        IconButtonCreator.getButton(TestSparkIcons.undoRemove, PluginLabelsBundle.get("undoRemoveTip"))
+
     // Create "Reset" button to reset the changes in the source code of the test
     private val resetButton = IconButtonCreator.getButton(TestSparkIcons.reset, PluginLabelsBundle.get("resetTip"))
 
@@ -135,6 +139,8 @@ class TestCasePanelBuilder(
     private val requestJLabel = JLabel(PluginLabelsBundle.get("requestJLabel"))
     private val requestComboBox = ComboBox(arrayOf("") + JsonEncoding.decode(llmSettingsState.defaultLLMRequests))
 
+    private val removedTestJLabel = JLabel(PluginLabelsBundle.get("removedTestJLabel"))
+
     private val sendButton = IconButtonCreator.getButton(TestSparkIcons.send, PluginLabelsBundle.get("send"))
 
     private val loadingLabel: JLabel = JLabel(TestSparkIcons.loading)
@@ -279,11 +285,17 @@ class TestCasePanelBuilder(
         loadingLabel.isVisible = false
         buttonsPanel.add(loadingLabel)
         buttonsPanel.add(Box.createHorizontalGlue())
+        removedTestJLabel.isVisible = false
+        buttonsPanel.add(removedTestJLabel)
+        buttonsPanel.add(Box.createHorizontalGlue())
         resetButton.isEnabled = false
         buttonsPanel.add(resetButton)
         resetToLastRunButton.isEnabled = false
         buttonsPanel.add(resetToLastRunButton)
         buttonsPanel.add(removeButton)
+        buttonsPanel.add(undoRemoveButton)
+        undoRemoveButton.isVisible = false
+
         buttonsPanel.add(Box.createRigidArea(Dimension(12, 0)))
 
         panel.add(requestPanel)
@@ -303,6 +315,7 @@ class TestCasePanelBuilder(
         resetButton.addActionListener { reset() }
         resetToLastRunButton.addActionListener { resetToLastRun() }
         removeButton.addActionListener { remove() }
+        undoRemoveButton.addActionListener { undoRemove() }
 
         sendButton.addActionListener { sendRequest() }
 
@@ -393,8 +406,8 @@ class TestCasePanelBuilder(
 
         currentCodes[currentRequestNumber - 1] = testCase.testCode
 
-        // select checkbox
-        checkbox.isSelected = true
+        // select checkbox if the test is not removed
+        checkbox.isSelected = !isRemoved
 
         if (testCaseCodeToListOfCoveredLines.containsKey(testCase.testCode)) {
             testCase.coveredLines = testCaseCodeToListOfCoveredLines[testCase.testCode]!!
@@ -469,6 +482,7 @@ class TestCasePanelBuilder(
         resetToLastRunButton.isEnabled = isEnabled
         resetButton.isEnabled = isEnabled
         removeButton.isEnabled = isEnabled
+        undoRemoveButton.isEnabled = isEnabled
         sendButton.isEnabled = isEnabled
     }
 
@@ -610,25 +624,68 @@ class TestCasePanelBuilder(
     }
 
     /**
-     * Removes the button listener for the test case.
+     * Marks the test as removed.
      *
-     * This method is responsible for:
-     * 1. Removing the highlighting of the test.
-     * 2. Removing the test case from the cache.
-     * 3. Updating the UI.
+     * 1. The state is propagated to other parts, including ReportUpdater.
+     * 2. The test is indicated as removed in the UI.
+     * 3. The test won't be runnable or selectable by any means until it's restored.
      */
     private fun remove() {
-        // Remove the test case from the cache
+        // Temporarily remove the test case
         GenerateTestsTabHelper.removeTestCase(testCase.testName, generatedTestsTabData)
 
         runTestButton.isEnabled = false
+        sendButton.isEnabled = false
+        resetButton.isEnabled = false
+        resetToLastRunButton.isEnabled = false
+
+        checkbox.isSelected = false
+        checkbox.isEnabled = false
+
+        removeButton.isVisible = false
+        undoRemoveButton.isVisible = true
+
+        languageTextField.isEnabled = false
+        removedTestJLabel.isVisible = true
+
         isRemoved = true
 
         ReportUpdater.removeTestCase(report, testCase, coverageVisualisationTabBuilder, generatedTestsTabData)
 
         GenerateTestsTabHelper.update(generatedTestsTabData)
     }
 
+    /**
+     * Marks the test as not removed.
+     *
+     * 1. The state is propagated to other parts, including ReportUpdater.
+     * 2. The test is indicated as normal in the UI.
+     * 3. The test is runnable and selectable like any other normal test.
+     */
+    private fun undoRemove() {
+        GenerateTestsTabHelper.undoRemoveTestCase(testCase.testName, generatedTestsTabData)
+
+        runTestButton.isEnabled = (getError() == null)
+        sendButton.isEnabled = true
+        resetButton.isEnabled = true
+        resetToLastRunButton.isEnabled = true
+
+        checkbox.isSelected = true
+        checkbox.isEnabled = true
+
+        removeButton.isVisible = true
+        undoRemoveButton.isVisible = false
+
+        languageTextField.isEnabled = true
+        removedTestJLabel.isVisible = false
+
+        isRemoved = false
+
+        ReportUpdater.addTestCase(report, testCase, coverageVisualisationTabBuilder, generatedTestsTabData)
+
+        GenerateTestsTabHelper.update(generatedTestsTabData)
+    }
+
     /**
      * Determines if the "Run" button is enabled.
      *

src/main/kotlin/org/jetbrains/research/testspark/display/generatedTests/TopButtonsPanelBuilder.kt

@@ -44,22 +44,18 @@ class TopButtonsPanelBuilder {
             .filter { !it.isRemoved() }
             .count { it.getError()?.isEmpty() == true }
 
-        val removedTestsCount = generatedTestsTabData.testCasePanelFactories.count { it.isRemoved() }
+        val nonRemovedTestsCount = generatedTestsTabData.testCaseNameToPanel.size - generatedTestsTabData.testsRemoved
 
-        if (generatedTestsTabData.testCasePanelFactories.size == removedTestsCount) {
-            removeAllButton.doClick()
-            return
-        }
         testsSelectedLabel.text = String.format(
             testsSelectedText,
             generatedTestsTabData.testsSelected,
-            generatedTestsTabData.testCaseNameToPanel.size,
+            nonRemovedTestsCount,
         )
         testsPassedLabel.text =
             String.format(
                 testsPassedText,
                 passedTestsCount,
-                generatedTestsTabData.testCaseNameToPanel.size,
+                nonRemovedTestsCount,
             )
         runAllButton.isEnabled = false
         for (testCasePanelFactory in generatedTestsTabData.testCasePanelFactories) {
@@ -71,20 +67,22 @@ class TopButtonsPanelBuilder {
 
     /**
      * Toggles check boxes so that they are either all selected or all not selected,
-     *  depending on the provided parameter.
+     *  depending on the provided parameter. This affects only non-removed tests.
      *
      *  @param selected whether the checkboxes have to be selected or not
      */
     private fun toggleAllCheckboxes(selected: Boolean, generatedTestsTabData: GeneratedTestsTabData) {
-        generatedTestsTabData.testCaseNameToPanel.forEach { (_, jPanel) ->
+        generatedTestsTabData.testCaseNameToPanel.forEach { (testCaseName, jPanel) ->
             val checkBox = jPanel.getComponent(0) as JCheckBox
-            checkBox.isSelected = selected
+            if (generatedTestsTabData.testCaseNameToEnabled[testCaseName]!!) {
+                checkBox.isSelected = selected
+            }
         }
         generatedTestsTabData.testsSelected = if (selected) generatedTestsTabData.testCaseNameToPanel.size else 0
     }
 
     /**
-     * Executes all test cases.
+     * Executes all non-removed test cases.
      *
      * This method presents a caution message to the user and asks for confirmation before executing the test cases.
      * If the user confirms, it iterates through each test case panel factory and runs the corresponding test.
@@ -106,7 +104,9 @@ class TopButtonsPanelBuilder {
         val tasks: Queue<(CustomProgressIndicator) -> Unit> = LinkedList()
 
         for (testCasePanelFactory in generatedTestsTabData.testCasePanelFactories) {
-            testCasePanelFactory.addTask(tasks)
+            if (!testCasePanelFactory.isRemoved()) {
+                testCasePanelFactory.addTask(tasks)
+            }
         }
         // run tasks one after each other
         executeTasks(project, tasks, generatedTestsTabData)

src/main/kotlin/org/jetbrains/research/testspark/display/utils/ReportUpdater.kt

@@ -29,6 +29,17 @@ object ReportUpdater {
         coverageVisualisationTabBuilder.show(report, generatedTestsTabData)
     }
 
+    fun addTestCase(
+        report: Report,
+        testCase: TestCase,
+        coverageVisualisationTabBuilder: CoverageVisualisationTabBuilder,
+        generatedTestsTabData: GeneratedTestsTabData,
+    ) {
+        report.testCaseList[testCase.id] = testCase
+        report.normalized()
+        coverageVisualisationTabBuilder.show(report, generatedTestsTabData)
+    }
+
     fun unselectTestCase(
         report: Report,
         unselectedTestCases: HashMap<Int, TestCase>,