From 6ba5ee3a834a0ba18f5596650bc03d30d94dc468 Mon Sep 17 00:00:00 2001
From: Yury Semikhatsky <yurys@chromium.org>
Date: Mon, 7 Oct 2024 15:42:12 -0700
Subject: [PATCH] chore(ui-mode): expand all button (#32994)

<img width="230" alt="image"
src="https://github.com/user-attachments/assets/dd7fa3a1-39ec-4b88-9279-d664c9c4e5cd">


Reference https://github.com/microsoft/playwright/issues/32825
---
 .../src/ui/uiModeTestListView.tsx             | 16 ++++++++++--
 packages/trace-viewer/src/ui/uiModeView.tsx   |  5 ++++
 .../playwright-test/ui-mode-test-tree.spec.ts | 26 +++++++++++++++++++
 3 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/packages/trace-viewer/src/ui/uiModeTestListView.tsx b/packages/trace-viewer/src/ui/uiModeTestListView.tsx
index 99a5f22dfa..ce1c0fef37 100644
--- a/packages/trace-viewer/src/ui/uiModeTestListView.tsx
+++ b/packages/trace-viewer/src/ui/uiModeTestListView.tsx
@@ -47,12 +47,14 @@ export const TestListView: React.FC<{
   isLoading?: boolean,
   onItemSelected: (item: { treeItem?: TreeItem, testCase?: reporterTypes.TestCase, testFile?: SourceLocation }) => void,
   requestedCollapseAllCount: number,
+  requestedExpandAllCount: number,
   setFilterText: (text: string) => void,
   onRevealSource: () => void,
-}> = ({ filterText, testModel, testServerConnection, testTree, runTests, runningState, watchAll, watchedTreeIds, setWatchedTreeIds, isLoading, onItemSelected, requestedCollapseAllCount, setFilterText, onRevealSource }) => {
+}> = ({ filterText, testModel, testServerConnection, testTree, runTests, runningState, watchAll, watchedTreeIds, setWatchedTreeIds, isLoading, onItemSelected, requestedCollapseAllCount, requestedExpandAllCount, setFilterText, onRevealSource }) => {
   const [treeState, setTreeState] = React.useState<TreeState>({ expandedItems: new Map() });
   const [selectedTreeItemId, setSelectedTreeItemId] = React.useState<string | undefined>();
   const [collapseAllCount, setCollapseAllCount] = React.useState(requestedCollapseAllCount);
+  const [expandAllCount, setExpandAllCount] = React.useState(requestedExpandAllCount);
 
   // Look for a first failure within the run batch to select it.
   React.useEffect(() => {
@@ -67,6 +69,16 @@ export const TestListView: React.FC<{
       return;
     }
 
+    if (expandAllCount !== requestedExpandAllCount) {
+      treeState.expandedItems.clear();
+      for (const item of testTree.flatTreeItems())
+        treeState.expandedItems.set(item.id, true);
+      setExpandAllCount(requestedExpandAllCount);
+      setSelectedTreeItemId(undefined);
+      setTreeState({ ...treeState });
+      return;
+    }
+
     if (!runningState || runningState.itemSelectedByUser)
       return;
     let selectedTreeItem: TreeItem | undefined;
@@ -85,7 +97,7 @@ export const TestListView: React.FC<{
 
     if (selectedTreeItem)
       setSelectedTreeItemId(selectedTreeItem.id);
-  }, [runningState, setSelectedTreeItemId, testTree, collapseAllCount, setCollapseAllCount, requestedCollapseAllCount, treeState, setTreeState]);
+  }, [runningState, setSelectedTreeItemId, testTree, collapseAllCount, setCollapseAllCount, requestedCollapseAllCount, expandAllCount, setExpandAllCount, requestedExpandAllCount, treeState, setTreeState]);
 
   // Compute selected item.
   const { selectedTreeItem } = React.useMemo(() => {
diff --git a/packages/trace-viewer/src/ui/uiModeView.tsx b/packages/trace-viewer/src/ui/uiModeView.tsx
index a5f1af7d99..69a5988641 100644
--- a/packages/trace-viewer/src/ui/uiModeView.tsx
+++ b/packages/trace-viewer/src/ui/uiModeView.tsx
@@ -90,6 +90,7 @@ export const UIModeView: React.FC<{}> = ({
   const commandQueue = React.useRef(Promise.resolve());
   const runTestBacklog = React.useRef<Set<string>>(new Set());
   const [collapseAllCount, setCollapseAllCount] = React.useState(0);
+  const [expandAllCount, setExpandAllCount] = React.useState(0);
   const [isDisconnected, setIsDisconnected] = React.useState(false);
   const [hasBrowsers, setHasBrowsers] = React.useState(true);
   const [testServerConnection, setTestServerConnection] = React.useState<TestServerConnection>();
@@ -473,6 +474,9 @@ export const UIModeView: React.FC<{}> = ({
           <ToolbarButton icon='collapse-all' title='Collapse all' onClick={() => {
             setCollapseAllCount(collapseAllCount + 1);
           }} />
+          <ToolbarButton icon='expand-all' title='Expand all' onClick={() => {
+            setExpandAllCount(expandAllCount + 1);
+          }} />
         </Toolbar>
         <TestListView
           filterText={filterText}
@@ -487,6 +491,7 @@ export const UIModeView: React.FC<{}> = ({
           setWatchedTreeIds={setWatchedTreeIds}
           isLoading={isLoading}
           requestedCollapseAllCount={collapseAllCount}
+          requestedExpandAllCount={expandAllCount}
           setFilterText={setFilterText}
           onRevealSource={onRevealSource}
         />
diff --git a/tests/playwright-test/ui-mode-test-tree.spec.ts b/tests/playwright-test/ui-mode-test-tree.spec.ts
index 7346c204e4..2c19ea0a87 100644
--- a/tests/playwright-test/ui-mode-test-tree.spec.ts
+++ b/tests/playwright-test/ui-mode-test-tree.spec.ts
@@ -297,6 +297,32 @@ test('should collapse all', async ({ runUITest }) => {
   `);
 });
 
+test('should expand all', {
+  annotation: { type: 'issue', description: 'https://github.com/microsoft/playwright/issues/32825' }
+}, async ({ runUITest }) => {
+  const { page } = await runUITest(basicTestTree);
+
+  await page.getByTestId('test-tree').getByText('suite').click();
+  await page.getByTitle('Collapse all').click();
+  await expect.poll(dumpTestTree(page)).toContain(`
+    ► ◯ a.test.ts
+    ► ◯ b.test.ts
+  `);
+
+  await page.getByTitle('Expand all').click();
+  await expect.poll(dumpTestTree(page)).toContain(`
+    ▼ ◯ a.test.ts
+        ◯ passes
+        ◯ fails
+      ▼ ◯ suite
+          ◯ inner passes
+          ◯ inner fails
+    ▼ ◯ b.test.ts
+        ◯ passes
+        ◯ fails
+  `);
+});
+
 test('should resolve title conflicts', async ({ runUITest }) => {
   const { page } = await runUITest({
     'a.test.ts': `