From af7e07e1510c4e039f9bd1273f7bade29e6a0ffd Mon Sep 17 00:00:00 2001 From: pyakovenko Date: Sun, 31 Aug 2025 19:07:55 +0200 Subject: [PATCH 1/4] feat(toolwindow): save state of the branches across files --- .../intellij/toolwindow/TestExplorerWindow.kt | 5 + .../intellij/toolwindow/TestFileTree.kt | 66 +++++++++++- .../plugin/intellij/toolwindow/treeutils.kt | 101 ++++++++++++++++++ 3 files changed, 167 insertions(+), 5 deletions(-) diff --git a/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestExplorerWindow.kt b/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestExplorerWindow.kt index 1aabe000..0c9615b6 100644 --- a/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestExplorerWindow.kt +++ b/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestExplorerWindow.kt @@ -104,6 +104,11 @@ class TestExplorerWindow(private val project: Project) : SimpleToolWindowPanel(t project.messageBus.connect().subscribe( FileEditorManagerListener.FILE_EDITOR_MANAGER, object : FileEditorManagerListener { + override fun fileClosed(source: FileEditorManager, file: com.intellij.openapi.vfs.VirtualFile) { + // when a file is closed, reset the one-time expanded state so reopening expands all again + tree.markFileClosed(file) + } + override fun selectionChanged(event: FileEditorManagerEvent) { val file = fileEditorManager.selectedEditor?.file if (file != null) { diff --git a/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestFileTree.kt b/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestFileTree.kt index 2ee4be15..1d9e001d 100644 --- a/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestFileTree.kt +++ b/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestFileTree.kt @@ -3,6 +3,7 @@ package io.kotest.plugin.intellij.toolwindow import com.intellij.ide.util.treeView.NodeRenderer import com.intellij.ide.util.treeView.PresentableNodeDescriptor import com.intellij.openapi.project.Project +import com.intellij.openapi.vfs.VirtualFile import com.intellij.ui.TreeUIHelper import javax.swing.tree.DefaultMutableTreeNode import javax.swing.tree.TreeModel @@ -17,6 +18,10 @@ class TestFileTree( private val kotestTestExplorerService: KotestTestExplorerService = project.getService(KotestTestExplorerService::class.java) private var initialized = false + private val filesInitiallyExpanded = mutableSetOf() + private var lastFileKey: String? = null + private val savedExpandedByFile = mutableMapOf>() + private val savedAllKeysByFile = mutableMapOf>() init { selectionModel.selectionMode = TreeSelectionModel.SINGLE_TREE_SELECTION @@ -40,13 +45,64 @@ class TestFileTree( super.setModel(treeModel) return } - val expanded = isExpanded(0) + val newFileKey = currentFileKey() + + // If switching away from a file, save its state first + if (lastFileKey != null && lastFileKey != newFileKey) { + savedExpandedByFile[lastFileKey!!] = collectExpandedPathKeys() + savedAllKeysByFile[lastFileKey!!] = collectAllPathKeys() + } + + val firstOpenForFile = newFileKey != null && !filesInitiallyExpanded.contains(newFileKey) + val sameFile = newFileKey == lastFileKey + + // Prepare previous keys/expansion baselines + val prevAllKeysForThisFile: Set = when { + firstOpenForFile -> emptySet() + sameFile -> collectAllPathKeys() + newFileKey != null -> savedAllKeysByFile[newFileKey] ?: emptySet() + else -> emptySet() + } + val expandedKeysToRestore: Set = when { + firstOpenForFile -> emptySet() + sameFile -> collectExpandedPathKeys() + newFileKey != null -> savedExpandedByFile[newFileKey] ?: emptySet() + else -> emptySet() + } + super.setModel(treeModel) - expandAllNodes() - setModuleGroupNodeExpandedState(expanded) + + // Compute added nodes relative to the previous snapshot of this file (if any) + val newAllKeys = collectAllPathKeys() + if (!firstOpenForFile) { + val addedKeys = newAllKeys - prevAllKeysForThisFile + if (addedKeys.isNotEmpty()) expandAncestorPrefixesFor(addedKeys) + } + + if (firstOpenForFile) { + // First time this file is shown in the tool window: expand everything + expandAllNodes() + newFileKey?.let { filesInitiallyExpanded.add(it) } + } else { + // Restore previous expansion state for this file + if (expandedKeysToRestore.isNotEmpty()) expandPathsByKeys(expandedKeysToRestore) + } + + // Update caches for this file and mark it as current + if (newFileKey != null) { + savedAllKeysByFile[newFileKey] = newAllKeys + savedExpandedByFile[newFileKey] = collectExpandedPathKeys() + } + lastFileKey = newFileKey } - private fun setModuleGroupNodeExpandedState(expanded: Boolean) { - if (expanded) expandRow(0) else collapseRow(0) + fun markFileClosed(file: VirtualFile) { + filesInitiallyExpanded.remove(file.path) + savedExpandedByFile.remove(file.path) + savedAllKeysByFile.remove(file.path) + if (lastFileKey == file.path) lastFileKey = null } + + private fun currentFileKey(): String? = kotestTestExplorerService.currentFile?.path + } diff --git a/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/treeutils.kt b/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/treeutils.kt index aa7fd53e..737d488a 100644 --- a/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/treeutils.kt +++ b/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/treeutils.kt @@ -34,3 +34,104 @@ fun TreePath.nodeDescriptor(): PresentableNodeDescriptor<*>? { else -> null } } + +// --- Expansion state helpers --- + +/** + * Builds a logical key for a node to preserve expansion state across model rebuilds. + * Only nodes that can have children are keyed (root, modules, module, tags, file, spec, container test). + */ +private fun DefaultMutableTreeNode.expansionKeyOrNull(): String? { + return when (val descriptor = userObject) { + is KotestRootNodeDescriptor -> "root" + is ModulesNodeDescriptor -> "modules" + is ModuleNodeDescriptor -> "module:${descriptor.module.name}" + is TagsNodeDescriptor -> "tags" + is TestFileNodeDescriptor -> "file" + is SpecNodeDescriptor -> "spec:${descriptor.fqn.asString()}" + is TestNodeDescriptor -> { + // Only container tests can have children; still safe to key all tests + "test:${descriptor.test.test.descriptorPath()}" + } + else -> null + } +} + +/** + * Returns a set of expansion path keys for the currently expanded nodes. + */ +fun JTree.collectExpandedPathKeys(): Set { + val root = model.root as? DefaultMutableTreeNode ?: return emptySet() + val expanded = mutableSetOf() + // Enumerate all expanded descendants starting from root + val enumeration = getExpandedDescendants(TreePath(root.path)) ?: return emptySet() + while (enumeration.hasMoreElements()) { + val path = enumeration.nextElement() + val key = pathToExpansionKey(path) + if (key != null) expanded.add(key) + } + return expanded +} + +/** + * Expands nodes in the current model whose logical expansion keys appear in [keys]. + */ +fun JTree.expandPathsByKeys(keys: Set) { + val root = model.root as? DefaultMutableTreeNode ?: return + fun recurse(node: DefaultMutableTreeNode, prefix: String?) { + val key = node.expansionKeyOrNull() + val pathKey = if (key == null) prefix else listOfNotNull(prefix, key).joinToString("/") + if (pathKey != null && keys.contains(pathKey)) { + expandPath(TreePath(node.path)) + } + val children = node.children() + while (children.hasMoreElements()) { + val child = children.nextElement() as DefaultMutableTreeNode + recurse(child, pathKey) + } + } + recurse(root, null) +} + +private fun pathToExpansionKey(path: TreePath): String? { + val parts = path.path + .mapNotNull { it as? DefaultMutableTreeNode } + .mapNotNull { it.expansionKeyOrNull() } + return if (parts.isEmpty()) null else parts.joinToString("/") +} + +/** + * Returns a set of path keys for all nodes in the current model. + */ +fun JTree.collectAllPathKeys(): Set { + val root = model.root as? DefaultMutableTreeNode ?: return emptySet() + val keys = mutableSetOf() + fun recurse(node: DefaultMutableTreeNode) { + val path = TreePath(node.path) + val key = pathToExpansionKey(path) + if (key != null) keys.add(key) + val children = node.children() + while (children.hasMoreElements()) { + recurse(children.nextElement() as DefaultMutableTreeNode) + } + } + recurse(root) + return keys +} + +/** + * Expands all ancestor prefixes for the given full path keys. + */ +fun JTree.expandAncestorPrefixesFor(keys: Set) { + if (keys.isEmpty()) return + val toExpand = mutableSetOf() + keys.forEach { key -> + val parts = key.split('/') + val acc = mutableListOf() + parts.forEach { part -> + acc.add(part) + toExpand.add(acc.joinToString("/")) + } + } + expandPathsByKeys(toExpand) +} From fa96a3578f5949f297abcf06af49e272eb215746 Mon Sep 17 00:00:00 2001 From: pyakovenko Date: Sun, 31 Aug 2025 19:24:30 +0200 Subject: [PATCH 2/4] feat(toolwindow): save state of the branches across files --- .../intellij/toolwindow/TestFileTree.kt | 42 ++++++++++--------- 1 file changed, 23 insertions(+), 19 deletions(-) diff --git a/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestFileTree.kt b/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestFileTree.kt index 1d9e001d..370111d2 100644 --- a/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestFileTree.kt +++ b/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestFileTree.kt @@ -9,6 +9,12 @@ import javax.swing.tree.DefaultMutableTreeNode import javax.swing.tree.TreeModel import javax.swing.tree.TreeSelectionModel +private data class FileTreeState( + val allKeys: Set, + val expandedKeys: Set, + var initiallyExpanded: Boolean, +) + class TestFileTree( project: Project, ) : com.intellij.ui.treeStructure.Tree(), @@ -18,10 +24,8 @@ class TestFileTree( private val kotestTestExplorerService: KotestTestExplorerService = project.getService(KotestTestExplorerService::class.java) private var initialized = false - private val filesInitiallyExpanded = mutableSetOf() private var lastFileKey: String? = null - private val savedExpandedByFile = mutableMapOf>() - private val savedAllKeysByFile = mutableMapOf>() + private val stateByFileKey = mutableMapOf() init { selectionModel.selectionMode = TreeSelectionModel.SINGLE_TREE_SELECTION @@ -49,24 +53,27 @@ class TestFileTree( // If switching away from a file, save its state first if (lastFileKey != null && lastFileKey != newFileKey) { - savedExpandedByFile[lastFileKey!!] = collectExpandedPathKeys() - savedAllKeysByFile[lastFileKey!!] = collectAllPathKeys() + val prevAll = collectAllPathKeys() + val prevExpanded = collectExpandedPathKeys() + val prevInit = stateByFileKey[lastFileKey!!]?.initiallyExpanded ?: false + stateByFileKey[lastFileKey!!] = FileTreeState(prevAll, prevExpanded, prevInit) } - val firstOpenForFile = newFileKey != null && !filesInitiallyExpanded.contains(newFileKey) val sameFile = newFileKey == lastFileKey + val prevStateForNew = if (newFileKey != null) stateByFileKey[newFileKey] else null + val firstOpenForFile = newFileKey != null && prevStateForNew == null - // Prepare previous keys/expansion baselines + // Baselines (use live tree for same file; fallback to stored state when switching) val prevAllKeysForThisFile: Set = when { firstOpenForFile -> emptySet() sameFile -> collectAllPathKeys() - newFileKey != null -> savedAllKeysByFile[newFileKey] ?: emptySet() + newFileKey != null -> prevStateForNew?.allKeys ?: emptySet() else -> emptySet() } val expandedKeysToRestore: Set = when { firstOpenForFile -> emptySet() sameFile -> collectExpandedPathKeys() - newFileKey != null -> savedExpandedByFile[newFileKey] ?: emptySet() + newFileKey != null -> prevStateForNew?.expandedKeys ?: emptySet() else -> emptySet() } @@ -80,26 +87,23 @@ class TestFileTree( } if (firstOpenForFile) { - // First time this file is shown in the tool window: expand everything + // First time this file is shown in the tool window: expand everything except Modules expandAllNodes() - newFileKey?.let { filesInitiallyExpanded.add(it) } + stateByFileKey[newFileKey] = FileTreeState(newAllKeys, collectExpandedPathKeys(), initiallyExpanded = true) } else { // Restore previous expansion state for this file if (expandedKeysToRestore.isNotEmpty()) expandPathsByKeys(expandedKeysToRestore) + if (newFileKey != null) { + val init = prevStateForNew?.initiallyExpanded ?: true + stateByFileKey[newFileKey] = FileTreeState(newAllKeys, collectExpandedPathKeys(), init) + } } - // Update caches for this file and mark it as current - if (newFileKey != null) { - savedAllKeysByFile[newFileKey] = newAllKeys - savedExpandedByFile[newFileKey] = collectExpandedPathKeys() - } lastFileKey = newFileKey } fun markFileClosed(file: VirtualFile) { - filesInitiallyExpanded.remove(file.path) - savedExpandedByFile.remove(file.path) - savedAllKeysByFile.remove(file.path) + stateByFileKey.remove(file.path) if (lastFileKey == file.path) lastFileKey = null } From 10a5b811b320eed215a8290d9ab021f1b655d2c9 Mon Sep 17 00:00:00 2001 From: pyakovenko Date: Sun, 31 Aug 2025 19:27:25 +0200 Subject: [PATCH 3/4] fix: import --- .../io/kotest/plugin/intellij/toolwindow/TestExplorerWindow.kt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestExplorerWindow.kt b/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestExplorerWindow.kt index 0c9615b6..8d9b4784 100644 --- a/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestExplorerWindow.kt +++ b/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestExplorerWindow.kt @@ -9,6 +9,7 @@ import com.intellij.openapi.fileEditor.FileEditorManagerEvent import com.intellij.openapi.fileEditor.FileEditorManagerListener import com.intellij.openapi.project.Project import com.intellij.openapi.ui.SimpleToolWindowPanel +import com.intellij.openapi.vfs.VirtualFile import com.intellij.openapi.vfs.VirtualFileManager import com.intellij.openapi.vfs.newvfs.BulkFileListener import com.intellij.openapi.vfs.newvfs.events.VFileEvent @@ -104,7 +105,7 @@ class TestExplorerWindow(private val project: Project) : SimpleToolWindowPanel(t project.messageBus.connect().subscribe( FileEditorManagerListener.FILE_EDITOR_MANAGER, object : FileEditorManagerListener { - override fun fileClosed(source: FileEditorManager, file: com.intellij.openapi.vfs.VirtualFile) { + override fun fileClosed(source: FileEditorManager, file: VirtualFile) { // when a file is closed, reset the one-time expanded state so reopening expands all again tree.markFileClosed(file) } From e025f11063e2096e68dfc219e59329fb03450d6d Mon Sep 17 00:00:00 2001 From: pyakovenko Date: Sun, 31 Aug 2025 19:28:36 +0200 Subject: [PATCH 4/4] fix: comment --- .../kotlin/io/kotest/plugin/intellij/toolwindow/TestFileTree.kt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestFileTree.kt b/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestFileTree.kt index 370111d2..3d8b756c 100644 --- a/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestFileTree.kt +++ b/src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestFileTree.kt @@ -87,7 +87,7 @@ class TestFileTree( } if (firstOpenForFile) { - // First time this file is shown in the tool window: expand everything except Modules + // First time this file is shown in the tool window: expand everything expandAllNodes() stateByFileKey[newFileKey] = FileTreeState(newAllKeys, collectExpandedPathKeys(), initiallyExpanded = true) } else {