Fix expand/collapse all buttons (#23590)

* communicate via customevents

* Handle open group logic in wrapper

* fix tests

* Make grid action buttons sticky

* Add default toggle fn

* fix splitting task id by '.'

* fix missing dagrun ids

(cherry picked from commit afdfece)

airflow/www/static/js/grid/AutoRefresh.jsx

@@ -0,0 +1,51 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import React from 'react';
+import {
+  Switch,
+  FormControl,
+  FormLabel,
+  Spinner,
+} from '@chakra-ui/react';
+
+import { useAutoRefresh } from './context/autorefresh';
+
+const AutoRefresh = () => {
+  const { isRefreshOn, toggleRefresh, isPaused } = useAutoRefresh();
+
+  return (
+    <FormControl display="flex" width="auto" mr={2}>
+      <Spinner color="blue.500" speed="1s" mr="4px" visibility={isRefreshOn ? 'visible' : 'hidden'} />
+      <FormLabel htmlFor="auto-refresh" mb={0} fontWeight="normal">
+        Auto-refresh
+      </FormLabel>
+      <Switch
+        id="auto-refresh"
+        onChange={() => toggleRefresh(true)}
+        isDisabled={isPaused}
+        isChecked={isRefreshOn}
+        size="lg"
+        title={isPaused ? 'Autorefresh is disabled while the DAG is paused' : ''}
+      />
+    </FormControl>
+  );
+};
+
+export default AutoRefresh;

airflow/www/static/js/grid/Grid.jsx

@@ -19,55 +19,39 @@
 
 /* global localStorage, ResizeObserver */
 
-import React, { useRef, useEffect } from 'react';
+import React, { useState, useRef, useEffect } from 'react';
 import {
   Table,
   Tbody,
   Box,
-  Switch,
-  FormControl,
-  FormLabel,
-  Spinner,
   Thead,
   Flex,
-  useDisclosure,
-  Button,
-  Divider,
 } from '@chakra-ui/react';
 
 import { useGridData } from './api';
 import renderTaskRows from './renderTaskRows';
 import ResetRoot from './ResetRoot';
 import DagRuns from './dagRuns';
-import Details from './details';
-import useSelection from './utils/useSelection';
-import { useAutoRefresh } from './context/autorefresh';
 import ToggleGroups from './ToggleGroups';
-import FilterBar from './FilterBar';
-import LegendRow from './LegendRow';
+import { getMetaValue } from '../utils';
+import AutoRefresh from './AutoRefresh';
 
-const sidePanelKey = 'hideSidePanel';
+const dagId = getMetaValue('dag_id');
 
-const Grid = () => {
+const Grid = ({ isPanelOpen = false }) => {
   const scrollRef = useRef();
   const tableRef = useRef();
 
   const { data: { groups, dagRuns } } = useGridData();
   const dagRunIds = dagRuns.map((dr) => dr.runId);
 
-  const { isRefreshOn, toggleRefresh, isPaused } = useAutoRefresh();
-  const isPanelOpen = localStorage.getItem(sidePanelKey) !== 'true';
-  const { isOpen, onToggle } = useDisclosure({ defaultIsOpen: isPanelOpen });
+  const openGroupsKey = `${dagId}/open-groups`;
+  const storedGroups = JSON.parse(localStorage.getItem(openGroupsKey)) || [];
+  const [openGroupIds, setOpenGroupIds] = useState(storedGroups);
 
-  const { clearSelection } = useSelection();
-  const toggleSidePanel = () => {
-    if (!isOpen) {
-      localStorage.setItem(sidePanelKey, false);
-    } else {
-      clearSelection();
-      localStorage.setItem(sidePanelKey, true);
-    }
-    onToggle();
+  const onToggleGroups = (groupIds) => {
+    localStorage.setItem(openGroupsKey, JSON.stringify(groupIds));
+    setOpenGroupIds(groupIds);
   };
 
   const scrollOnResize = new ResizeObserver(() => {
@@ -91,64 +75,42 @@ const Grid = () => {
   }, [tableRef, scrollOnResize]);
 
   return (
-    <Box mt={3}>
-      <FilterBar />
-      <LegendRow />
-      <Divider mb={5} borderBottomWidth={2} />
-      <Flex flexGrow={1} justifyContent="space-between" alignItems="center">
-        <Flex alignItems="center">
-          <FormControl display="flex" width="auto" mr={2}>
-            {isRefreshOn && <Spinner color="blue.500" speed="1s" mr="4px" />}
-            <FormLabel htmlFor="auto-refresh" mb={0} fontWeight="normal">
-              Auto-refresh
-            </FormLabel>
-            <Switch
-              id="auto-refresh"
-              onChange={() => toggleRefresh(true)}
-              isDisabled={isPaused}
-              isChecked={isRefreshOn}
-              size="lg"
-              title={isPaused ? 'Autorefresh is disabled while the DAG is paused' : ''}
-            />
-          </FormControl>
-          <ToggleGroups groups={groups} />
-          <ResetRoot />
-        </Flex>
-        <Button
-          onClick={toggleSidePanel}
-          aria-label={isOpen ? 'Show Details' : 'Hide Details'}
-          variant={isOpen ? 'solid' : 'outline'}
-        >
-          {isOpen ? 'Hide ' : 'Show '}
-          Details Panel
-        </Button>
+    <Box
+      position="relative"
+      m={3}
+      mt={0}
+      overflow="auto"
+      ref={scrollRef}
+      flexGrow={1}
+      minWidth={isPanelOpen && '300px'}
+    >
+      <Flex alignItems="center" position="sticky" top={0} left={0}>
+        <AutoRefresh />
+        <ToggleGroups
+          groups={groups}
+          openGroupIds={openGroupIds}
+          onToggleGroups={onToggleGroups}
+        />
+        <ResetRoot />
       </Flex>
-      <Flex flexDirection="row" justifyContent="space-between">
-        <Box
-          position="relative"
-          mt={2}
-          m="12px"
-          overflow="auto"
-          ref={scrollRef}
-          flexGrow={1}
-          minWidth={isOpen && '300px'}
+      <Table>
+        <Thead display="block" pr="10px" position="sticky" top={0} zIndex={2} bg="white">
+          <DagRuns />
+        </Thead>
+        {/* TODO: remove hardcoded values. 665px is roughly the total heade+footer height */}
+        <Tbody
+          display="block"
+          width="100%"
+          maxHeight="calc(100vh - 665px)"
+          minHeight="500px"
+          ref={tableRef}
+          pr="10px"
         >
-          <Table>
-            <Thead display="block" pr="10px" position="sticky" top={0} zIndex={2} bg="white">
-              <DagRuns />
-            </Thead>
-            {/* TODO: remove hardcoded values. 665px is roughly the total header+footer height */}
-            <Tbody display="block" width="100%" maxHeight="calc(100vh - 665px)" minHeight="500px" ref={tableRef} pr="10px">
-              {renderTaskRows({
-                task: groups, dagRunIds,
-              })}
-            </Tbody>
-          </Table>
-        </Box>
-        {isOpen && (
-          <Details />
-        )}
-      </Flex>
+          {renderTaskRows({
+            task: groups, dagRunIds, openGroupIds, onToggleGroups,
+          })}
+        </Tbody>
+      </Table>
     </Box>
   );
 };

airflow/www/static/js/grid/ToggleGroups.test.jsx → airflow/www/static/js/grid/Grid.test.jsx

@@ -1,3 +1,4 @@
+/* eslint-disable class-methods-use-this */
 /*!
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
@@ -17,15 +18,14 @@
  * under the License.
  */
 
-/* global describe, test, expect */
+/* global describe, test, expect, beforeEach, beforeAll, jest, window */
 
 import React from 'react';
-import { Flex, Table, Tbody } from '@chakra-ui/react';
 import { render, fireEvent, waitFor } from '@testing-library/react';
 
-import renderTaskRows from './renderTaskRows';
-import ToggleGroups from './ToggleGroups';
+import Grid from './Grid';
 import { Wrapper } from './utils/testUtils';
+import * as useGridDataModule from './api/useGridData';
 
 const mockGridData = {
   groups: {
@@ -115,9 +115,40 @@ const EXPAND = 'Expand all task groups';
 const COLLAPSE = 'Collapse all task groups';
 
 describe('Test ToggleGroups', () => {
+  beforeAll(() => {
+    class ResizeObserver {
+      observe() {}
+
+      unobserve() {}
+
+      disconnect() {}
+    }
+
+    window.ResizeObserver = ResizeObserver;
+  });
+
+  beforeEach(() => {
+    jest.spyOn(useGridDataModule, 'default').mockImplementation(() => ({
+      data: mockGridData,
+    }));
+  });
+
+  test('Group defaults to closed', () => {
+    const { getByTestId, getByText, getAllByTestId } = render(
+      <Grid />,
+      { wrapper: Wrapper },
+    );
+
+    const groupName = getByText('group_1');
+
+    expect(getAllByTestId('task-instance')).toHaveLength(1);
+    expect(groupName).toBeInTheDocument();
+    expect(getByTestId('closed-group')).toBeInTheDocument();
+  });
+
   test('Buttons are disabled if all groups are expanded or collapsed', () => {
     const { getByTitle } = render(
-      <ToggleGroups groups={mockGridData.groups} />,
+      <Grid />,
       { wrapper: Wrapper },
     );
 
@@ -134,19 +165,8 @@ describe('Test ToggleGroups', () => {
   });
 
   test('Expand/collapse buttons toggle nested groups', async () => {
-    global.gridData = mockGridData;
-    const dagRunIds = mockGridData.dagRuns.map((dr) => dr.runId);
-    const task = mockGridData.groups;
-
     const { getByText, queryAllByTestId, getByTitle } = render(
-      <Flex>
-        <ToggleGroups groups={task} />
-        <Table>
-          <Tbody>
-            {renderTaskRows({ task, dagRunIds })}
-          </Tbody>
-        </Table>
-      </Flex>,
+      <Grid />,
       { wrapper: Wrapper },
     );
 

airflow/www/static/js/grid/Main.jsx

@@ -0,0 +1,80 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/* global localStorage */
+
+import React from 'react';
+import {
+  Box,
+  Flex,
+  useDisclosure,
+  Button,
+  Divider,
+} from '@chakra-ui/react';
+
+import Details from './details';
+import useSelection from './utils/useSelection';
+import Grid from './Grid';
+import FilterBar from './FilterBar';
+import LegendRow from './LegendRow';
+
+const detailsPanelKey = 'hideDetailsPanel';
+
+const Main = () => {
+  const isPanelOpen = localStorage.getItem(detailsPanelKey) !== 'true';
+  const { isOpen, onToggle } = useDisclosure({ defaultIsOpen: isPanelOpen });
+  const { clearSelection } = useSelection();
+
+  const toggleDetailsPanel = () => {
+    if (!isOpen) {
+      localStorage.setItem(detailsPanelKey, false);
+    } else {
+      clearSelection();
+      localStorage.setItem(detailsPanelKey, true);
+    }
+    onToggle();
+  };
+
+  return (
+    <Box>
+      <FilterBar />
+      <LegendRow />
+      <Divider mb={5} borderBottomWidth={2} />
+      <Flex flexDirection="row" justifyContent="space-between">
+        <Grid isPanelOpen={isOpen} />
+        <Box borderLeftWidth={isOpen ? 1 : 0} position="relative">
+          <Button
+            position="absolute"
+            top={0}
+            right={0}
+            onClick={toggleDetailsPanel}
+            aria-label={isOpen ? 'Show Details' : 'Hide Details'}
+            variant={isOpen ? 'solid' : 'outline'}
+          >
+            {isOpen ? 'Hide ' : 'Show '}
+            Details Panel
+          </Button>
+          {isOpen && (<Details />)}
+        </Box>
+      </Flex>
+    </Box>
+  );
+};
+
+export default Main;