From c0c5a2039e0eccab1d1f80bea2ccf7b01e5e4ddf Mon Sep 17 00:00:00 2001
From: Huong Nguyen <huong.nguyen@mckinsey.com>
Date: Tue, 14 Jan 2025 23:18:21 +0700
Subject: [PATCH] Add new reducer to hanlde update state for when toggle
 expanding all pipelines

Signed-off-by: Huong Nguyen <huong.nguyen@mckinsey.com>
---
 src/actions/modular-pipelines.js              |  9 ++++++
 .../flowchart-primary-toolbar.js              | 12 +++----
 src/reducers/modular-pipelines.js             | 32 +++++++++++++++++++
 3 files changed, 46 insertions(+), 7 deletions(-)

diff --git a/src/actions/modular-pipelines.js b/src/actions/modular-pipelines.js
index add0e75eb9..217a954277 100644
--- a/src/actions/modular-pipelines.js
+++ b/src/actions/modular-pipelines.js
@@ -55,3 +55,12 @@ export function toggleModularPipelinesExpanded(expandedIDs) {
     expandedIDs,
   };
 }
+
+export const TOGGLE_ALL_MODULAR_PIPELINES_EXPANDED =
+  'TOGGLE_ALL_MODULAR_PIPELINES_EXPANDED';
+export function toggleAllModularPipelinesExpanded(expandAllPipelines) {
+  return {
+    type: TOGGLE_ALL_MODULAR_PIPELINES_EXPANDED,
+    expandAllPipelines,
+  };
+}
diff --git a/src/components/flowchart-primary-toolbar/flowchart-primary-toolbar.js b/src/components/flowchart-primary-toolbar/flowchart-primary-toolbar.js
index 88de1bb5e9..7ab383aca9 100644
--- a/src/components/flowchart-primary-toolbar/flowchart-primary-toolbar.js
+++ b/src/components/flowchart-primary-toolbar/flowchart-primary-toolbar.js
@@ -7,7 +7,10 @@ import {
   toggleTextLabels,
   toggleExpandAllPipelines,
 } from '../../actions';
-import { toggleModularPipelinesExpanded } from '../../actions/modular-pipelines';
+import {
+  toggleModularPipelinesExpanded,
+  toggleAllModularPipelinesExpanded,
+} from '../../actions/modular-pipelines';
 import IconButton from '../ui/icon-button';
 import LabelIcon from '../icons/label';
 import ExportIcon from '../icons/export';
@@ -42,12 +45,6 @@ export const FlowchartPrimaryToolbar = ({
 
   const handleToggleExpandAllPipelines = () => {
     const isExpanded = !expandedPipelines;
-    // Exclude root from modularPipelineIDs
-    const filteredModularPipelineIDs = modularPipelineIDs.filter(
-      (id) => id !== '__root__'
-    );
-    // Pass an empty array when collapsing all pipelines
-    onToggleExpandPipelines(isExpanded ? filteredModularPipelineIDs : []);
     onToggleExpandAllPipelines(isExpanded);
     toSetQueryParam('expandAllPipelines', isExpanded.toString());
   };
@@ -138,6 +135,7 @@ export const mapDispatchToProps = (dispatch) => ({
   },
   onToggleExpandAllPipelines: (isExpanded) => {
     dispatch(toggleExpandAllPipelines(isExpanded));
+    dispatch(toggleAllModularPipelinesExpanded(isExpanded));
   },
 });
 
diff --git a/src/reducers/modular-pipelines.js b/src/reducers/modular-pipelines.js
index 8698eace98..d42c98de38 100644
--- a/src/reducers/modular-pipelines.js
+++ b/src/reducers/modular-pipelines.js
@@ -3,6 +3,7 @@ import {
   TOGGLE_MODULAR_PIPELINES_EXPANDED,
   TOGGLE_SINGLE_MODULAR_PIPELINE_EXPANDED,
   TOGGLE_MODULAR_PIPELINE_DISABLED,
+  TOGGLE_ALL_MODULAR_PIPELINES_EXPANDED,
 } from '../actions/modular-pipelines';
 
 function modularPipelineReducer(modularPipelineState = {}, action) {
@@ -110,6 +111,37 @@ function modularPipelineReducer(modularPipelineState = {}, action) {
         visible: newVisibleState,
       });
     }
+    case TOGGLE_ALL_MODULAR_PIPELINES_EXPANDED: {
+      let newVisibleState = {};
+
+      // Determine which IDs should be expanded based on the action
+      const expandedIDs = action.expandAllPipelines
+        ? modularPipelineState.ids
+        : [];
+
+      if (action.expandAllPipelines) {
+        // If expanding all pipelines, set visibility for each pipeline and its children
+        modularPipelineState.ids.forEach((modularPipelineID) => {
+          newVisibleState[modularPipelineID] = false;
+          modularPipelineState.tree[modularPipelineID].children.forEach(
+            (child) => (newVisibleState[child.id] = true)
+          );
+        });
+      } else {
+        // If not expanding all, only set visibility for the children of the root pipeline
+        if (modularPipelineState.tree['__root__']) {
+          for (const child of modularPipelineState.tree['__root__'].children ||
+            []) {
+            newVisibleState[child.id] = true;
+          }
+        }
+      }
+
+      return updateState({
+        expanded: expandedIDs,
+        visible: newVisibleState,
+      });
+    }
 
     default:
       return modularPipelineState;