Add ability to restore test case if it was deleted

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

@@ -7,14 +7,6 @@ object GenerateTestsTabHelper {
      * @param testCaseName the name of the test
      */
     fun removeTestCase(testCaseName: String, generatedTestsTabData: GeneratedTestsTabData) {
-        // Update the number of selected test cases if necessary
-        if (generatedTestsTabData.testCaseNameToSelectedCheckbox[testCaseName]!!.isSelected) {
-            generatedTestsTabData.testsSelected--
-        }
-
-        // Remove the test panel from the UI
-        generatedTestsTabData.allTestCasePanel.remove(generatedTestsTabData.testCaseNameToPanel[testCaseName])
-
         // Remove the test panel
         generatedTestsTabData.testCaseNameToPanel.remove(testCaseName)
 
@@ -23,6 +15,22 @@ object GenerateTestsTabHelper {
 
         // Remove the editorTextField
         generatedTestsTabData.testCaseNameToEditorTextField.remove(testCaseName)
+
+        // Remove the Panel for
+        generatedTestsTabData.testCaseNameToRemovePanel.remove(testCaseName)
+    }
+
+    /**
+     * A helper method to remove a test case from the UI.
+     *
+     * @param testCaseName the name of the test
+     */
+    fun removeUITestCase(testCaseName: String, generatedTestsTabData: GeneratedTestsTabData) {
+        // Change test from selected to unselected
+        generatedTestsTabData.testCaseNameToSelectedCheckbox[testCaseName]!!.setSelected(false)
+
+        // Remove the test panel from the UI
+        generatedTestsTabData.allTestCasePanel.remove(generatedTestsTabData.testCaseNameToPanel[testCaseName])
     }
 
     /**

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

@@ -166,10 +166,16 @@ class GeneratedTestsTabBuilder(
             generatedTestsTabData.allTestCasePanel.add(testCasePanel)
             addSeparator()
 
+            val testPanelIndex: Int =
+                generatedTestsTabData.allTestCasePanel.getComponentZOrder(testCasePanel)
+
+            val testCaseRemovePanel: JPanel = testCasePanelBuilder.getRemovePanel(testPanelIndex)
+
             generatedTestsTabData.testCaseNameToPanel[testCase.testName] = testCasePanel
             generatedTestsTabData.testCaseNameToSelectedCheckbox[testCase.testName] = checkbox
             generatedTestsTabData.testCaseNameToEditorTextField[testCase.testName] =
                 testCasePanelBuilder.getEditorTextField()
+            generatedTestsTabData.testCaseNameToRemovePanel[testCase.testName] = testCaseRemovePanel
         }
         generatedTestsTabData.testsSelected = generatedTestsTabData.testCaseNameToPanel.size
         generatedTestsTabData.testCasePanelFactories.addAll(testCasePanelFactories)

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

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

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

@@ -48,6 +48,7 @@ import org.jetbrains.research.testspark.tools.ToolUtils
 import org.jetbrains.research.testspark.tools.factories.TestCompilerFactory
 import org.jetbrains.research.testspark.tools.llm.error.LLMErrorManager
 import org.jetbrains.research.testspark.tools.llm.test.JUnitTestSuitePresenter
+import java.awt.BorderLayout
 import java.awt.Dimension
 import java.awt.Toolkit
 import java.awt.datatransfer.Clipboard
@@ -302,7 +303,7 @@ class TestCasePanelBuilder(
         }
         resetButton.addActionListener { reset() }
         resetToLastRunButton.addActionListener { resetToLastRun() }
-        removeButton.addActionListener { remove() }
+        removeButton.addActionListener { switchTestPanel() }
 
         sendButton.addActionListener { sendRequest() }
 
@@ -316,6 +317,57 @@ class TestCasePanelBuilder(
         return panel
     }
 
+    /**
+     * Create a Panel that will replace test that was deleted by user.
+     */
+    fun getRemovePanel(testPanelIndex: Int): JPanel {
+        val panel = JPanel()
+
+        panel.layout = BoxLayout(panel, BoxLayout.X_AXIS)
+
+        val returnTestButton = JButton("Return test")
+        val removeTestButton = JButton("Remove test")
+
+        returnTestButton.isOpaque = false
+        returnTestButton.isContentAreaFilled = false
+        returnTestButton.isBorderPainted = true
+
+        removeTestButton.isOpaque = false
+        removeTestButton.isContentAreaFilled = false
+        removeTestButton.isBorderPainted = true
+
+        returnTestButton.addActionListener { undoRemove(testPanelIndex) }
+
+        removeTestButton.addActionListener { remove() }
+
+        panel.add(Box.createHorizontalGlue())
+
+        panel.add(returnTestButton)
+
+        panel.add(removeTestButton)
+
+        return panel
+    }
+
+    /**
+     * Method that switch removed panel to test panel that was deleted.
+     */
+    private fun undoRemove(testPanelIndex: Int) {
+        generatedTestsTabData.allTestCasePanel.remove(generatedTestsTabData.testCaseNameToRemovePanel[testCase.testName])
+
+        runTestButton.isEnabled = true
+        isRemoved = false
+
+        generatedTestsTabData.testCaseNameToSelectedCheckbox[testCase.testName]!!.setSelected(true)
+
+        generatedTestsTabData.allTestCasePanel.add(
+            generatedTestsTabData.testCaseNameToPanel[testCase.testName],
+            testPanelIndex
+        )
+
+        GenerateTestsTabHelper.update(generatedTestsTabData)
+    }
+
     /**
      * Updates the label displaying the request number information.
      * Uses the requestNumber template to format the label text.
@@ -618,6 +670,8 @@ class TestCasePanelBuilder(
      * 3. Updating the UI.
      */
     private fun remove() {
+        generatedTestsTabData.allTestCasePanel.remove(generatedTestsTabData.testCaseNameToRemovePanel[testCase.testName])
+
         // Remove the test case from the cache
         GenerateTestsTabHelper.removeTestCase(testCase.testName, generatedTestsTabData)
 
@@ -629,6 +683,27 @@ class TestCasePanelBuilder(
         GenerateTestsTabHelper.update(generatedTestsTabData)
     }
 
+    /**
+     * Switch test panel to removed panel of this test.
+     */
+    private fun switchTestPanel() {
+        val index: Int =
+            generatedTestsTabData.allTestCasePanel.getComponentZOrder(generatedTestsTabData.testCaseNameToPanel[testCase.testName])
+
+        GenerateTestsTabHelper.removeUITestCase(testCase.testName, generatedTestsTabData)
+
+        runTestButton.isEnabled = false
+        isRemoved = true
+
+        generatedTestsTabData.allTestCasePanel.add(
+            generatedTestsTabData.testCaseNameToRemovePanel[testCase.testName],
+            BorderLayout.EAST,
+            index
+        )
+
+        GenerateTestsTabHelper.update(generatedTestsTabData)
+    }
+
     /**
      * Determines if the "Run" button is enabled.
      *
@@ -698,7 +773,8 @@ class TestCasePanelBuilder(
      * Updates the current test case with the specified test name and test code.
      */
     private fun updateTestCaseInformation() {
-        testCase.testName = TestAnalyzerFactory.create(language).extractFirstTestMethodName(testCase.testName, languageTextField.document.text)
+        testCase.testName = TestAnalyzerFactory.create(language)
+            .extractFirstTestMethodName(testCase.testName, languageTextField.document.text)
         testCase.testCode = languageTextField.document.text
     }