Feature: Undo test case deletion

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

@@ -2,27 +2,53 @@ 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 remove a test case from the cache.
      *
      * @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)
+        val panel = generatedTestsTabData.testCaseNameToPanel.remove(testCaseName)
 
         // Remove the selected checkbox
-        generatedTestsTabData.testCaseNameToSelectedCheckbox.remove(testCaseName)
+        val checkbox = generatedTestsTabData.testCaseNameToSelectedCheckbox.remove(testCaseName)
 
         // Remove the editorTextField
-        generatedTestsTabData.testCaseNameToEditorTextField.remove(testCaseName)
+        val editorTextField = generatedTestsTabData.testCaseNameToEditorTextField.remove(testCaseName)
+
+        // Save data for possible restoration
+        val deletedTest = DeletedTest(panel!!, checkbox!!, editorTextField!!)
+        generatedTestsTabData.testCaseNameToDeletedTestData[testCaseName] = deletedTest
+    }
+
+    /**
+     * A helper method to restore a previously deleted test case to the cache.
+     *
+     * @param testCaseName the name of the test
+     */
+    fun restoreTestCase(testCaseName: String, generatedTestsTabData: GeneratedTestsTabData) {
+        // Retrieve information about previously deleted test
+        val deletedTest = generatedTestsTabData.testCaseNameToDeletedTestData.remove(testCaseName)!!
+
+        // Re-add the test panel
+        generatedTestsTabData.testCaseNameToPanel[testCaseName] = deletedTest.panel
+
+        // Re-add the selected checkbox
+        generatedTestsTabData.testCaseNameToSelectedCheckbox[testCaseName] = deletedTest.checkbox
+
+        // Re-add the editorTextField
+        generatedTestsTabData.testCaseNameToEditorTextField[testCaseName] = deletedTest.editorTextField
+    }
+
+    /**
+     * A helper method to remove all test cases from the cache and UI.
+     */
+    fun clear(generatedTestsTabData: GeneratedTestsTabData) {
+        generatedTestsTabData.allTestCasePanel.removeAll()
+        generatedTestsTabData.testCaseNameToPanel.clear()
+        generatedTestsTabData.testCaseNameToSelectedCheckbox.clear()
+        generatedTestsTabData.testCaseNameToEditorTextField.clear()
+        generatedTestsTabData.testCaseNameToDeletedTestData.clear()
     }
 
     /**

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

@@ -245,8 +245,7 @@ class GeneratedTestsTabBuilder(
      * Clears all the generated test cases from the UI and the internal cache.
      */
     fun clear() {
-        generatedTestsTabData.testCaseNameToPanel.toMap()
-            .forEach { GenerateTestsTabHelper.removeTestCase(it.key, generatedTestsTabData) }
+        GenerateTestsTabHelper.clear(generatedTestsTabData)
         generatedTestsTabData.testCasePanelFactories.clear()
         generatedTestsTabData.topButtonsPanelBuilder.clear(generatedTestsTabData)
 

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

@@ -14,6 +14,7 @@ class GeneratedTestsTabData {
     val testCaseNameToPanel: HashMap<String, JPanel> = HashMap()
     val testCaseNameToSelectedCheckbox: HashMap<String, JCheckBox> = HashMap()
     val testCaseNameToEditorTextField: HashMap<String, EditorTextField> = HashMap()
+    val testCaseNameToDeletedTestData: HashMap<String, DeletedTest> = HashMap()
     var testsSelected: Int = 0
     val unselectedTestCases: HashMap<Int, TestCase> = HashMap()
     val testCasePanelFactories: ArrayList<TestCasePanelBuilder> = arrayListOf()
@@ -28,3 +29,5 @@ class GeneratedTestsTabData {
     var contentManager: ContentManager? = null
     var content: Content? = null
 }
+
+data class DeletedTest(val panel: JPanel, val checkbox: JCheckBox, val editorTextField: EditorTextField)

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

@@ -52,6 +52,7 @@ import java.awt.Dimension
 import java.awt.Toolkit
 import java.awt.datatransfer.Clipboard
 import java.awt.datatransfer.StringSelection
+import java.awt.event.MouseAdapter
 import java.util.Queue
 import javax.swing.Box
 import javax.swing.BoxLayout
@@ -60,6 +61,7 @@ import javax.swing.JCheckBox
 import javax.swing.JLabel
 import javax.swing.JOptionPane
 import javax.swing.JPanel
+import javax.swing.OverlayLayout
 import javax.swing.ScrollPaneConstants
 import javax.swing.SwingUtilities
 import javax.swing.border.Border
@@ -302,7 +304,12 @@ class TestCasePanelBuilder(
         }
         resetButton.addActionListener { reset() }
         resetToLastRunButton.addActionListener { resetToLastRun() }
-        removeButton.addActionListener { remove() }
+
+        val deletedTestPanel = createDeletedTestPanel()
+        removeButton.addActionListener {
+            deletedTestPanel.isVisible = true
+            remove()
+        }
 
         sendButton.addActionListener { sendRequest() }
 
@@ -313,7 +320,40 @@ class TestCasePanelBuilder(
         requestComboBox.preferredSize = Dimension(0, 0)
         requestComboBox.isEditable = true
 
-        return panel
+        val overlayPanel = object : JPanel() {
+            override fun isOptimizedDrawingEnabled() = false
+        }.apply {
+            layout = OverlayLayout(this)
+            add(deletedTestPanel)
+            add(panel)
+        }
+
+        return overlayPanel
+    }
+
+    /**
+     * Creates the panel that will be displayed as the bottom panel for a deleted test case.
+     */
+    private fun createDeletedTestPanel(): JPanel {
+        val deletedTestPanel = JPanel().apply {
+            this.isVisible = false
+            this.addMouseListener(object : MouseAdapter() {})
+        }
+
+        deletedTestPanel.add(Box.createRigidArea(Dimension(15, 0)))
+        deletedTestPanel.add(JLabel("This test case has been deleted."))
+        deletedTestPanel.add(Box.createHorizontalGlue())
+
+        val undoButton = JButton("Undo")
+        deletedTestPanel.add(undoButton)
+        deletedTestPanel.add(Box.createRigidArea(Dimension(12, 0)))
+
+        undoButton.addActionListener {
+            deletedTestPanel.isVisible = false
+            restore()
+        }
+
+        return deletedTestPanel
     }
 
     /**
@@ -630,14 +670,38 @@ class TestCasePanelBuilder(
         // Remove the test case from the cache
         GenerateTestsTabHelper.removeTestCase(testCase.testName, generatedTestsTabData)
 
+        // Disable UI
+        enableLocalComponents(false)
+        languageTextField.isEnabled = false
+        checkbox.isSelected = false
+        checkbox.isEnabled = false
         runTestButton.isEnabled = false
+
         isRemoved = true
 
         ReportUpdater.removeTestCase(report, testCase, coverageVisualisationTabBuilder, generatedTestsTabData)
 
         GenerateTestsTabHelper.update(generatedTestsTabData)
     }
 
+    private fun restore() {
+        // Restore test case to cache
+        GenerateTestsTabHelper.restoreTestCase(testCase.testName, generatedTestsTabData)
+
+        // Enable UI
+        enableLocalComponents(true)
+        languageTextField.isEnabled = true
+        checkbox.isSelected = true
+        checkbox.isEnabled = true
+        runTestButton.isEnabled = true
+
+        isRemoved = false
+
+        ReportUpdater.restoreTestCase(report, testCase, coverageVisualisationTabBuilder, generatedTestsTabData)
+
+        GenerateTestsTabHelper.update(generatedTestsTabData)
+    }
+
     /**
      * Determines if the "Run" button is enabled.
      *

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

@@ -29,6 +29,17 @@ object ReportUpdater {
         coverageVisualisationTabBuilder.show(report, generatedTestsTabData)
     }
 
+    fun restoreTestCase(
+        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>,