Do not render section if it is not open (#5025)

* Do not render section if not opened

* Reset state when closing a plot section

Author: sroy3

extension/src/plots/webview/messages.ts

@@ -140,11 +140,26 @@ export class WebviewMessages {
           { isImage: !!message.payload },
           undefined
         )
+      case MessageFromWebviewType.REFRESH_PLOTS:
+        return this.sendDataForSection(message.payload)
       default:
         Logger.error(`Unexpected message: ${JSON.stringify(message)}`)
     }
   }
 
+  private sendDataForSection(section: PlotsSection) {
+    switch (section) {
+      case PlotsSection.COMPARISON_TABLE:
+        return this.sendComparisonPlots()
+      case PlotsSection.CUSTOM_PLOTS:
+        return this.sendCustomPlots()
+      case PlotsSection.TEMPLATE_PLOTS:
+        return this.sendTemplatePlots()
+      default:
+        return this.sendWebviewMessage()
+    }
+  }
+
   private async getWebviewData(): Promise<TPlotsData> {
     const selectedRevisions = this.plots.getSelectedRevisionDetails()
 

extension/src/test/suite/plots/index.test.ts

@@ -533,6 +533,60 @@ suite('Plots Test Suite', () => {
       )
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
+    it('should handle refresh plots message from the webview for the template plots', async () => {
+      const { mockMessageReceived, messageSpy } = await buildPlotsWebview({
+        disposer: disposable,
+        plotsDiff: plotsDiffFixture
+      })
+
+      messageSpy.resetHistory()
+      mockMessageReceived.fire({
+        payload: PlotsSection.TEMPLATE_PLOTS,
+        type: MessageFromWebviewType.REFRESH_PLOTS
+      })
+
+      expect(messageSpy).to.be.calledOnce
+      expect(messageSpy).to.be.calledWithExactly({
+        template: templatePlotsFixture
+      })
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
+
+    it('should handle refresh plots message from the webview for the custom plots', async () => {
+      const { mockMessageReceived, messageSpy } = await buildPlotsWebview({
+        disposer: disposable,
+        plotsDiff: plotsDiffFixture
+      })
+
+      messageSpy.resetHistory()
+      mockMessageReceived.fire({
+        payload: PlotsSection.CUSTOM_PLOTS,
+        type: MessageFromWebviewType.REFRESH_PLOTS
+      })
+
+      expect(messageSpy).to.be.calledOnce
+      expect(messageSpy).to.be.calledWithExactly({
+        custom: customPlotsFixture
+      })
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
+
+    it('should handle refresh plots message from the webview for the comparison table', async () => {
+      const { mockMessageReceived, messageSpy } = await buildPlotsWebview({
+        disposer: disposable,
+        plotsDiff: plotsDiffFixture
+      })
+
+      messageSpy.resetHistory()
+      mockMessageReceived.fire({
+        payload: PlotsSection.COMPARISON_TABLE,
+        type: MessageFromWebviewType.REFRESH_PLOTS
+      })
+
+      expect(messageSpy).to.be.calledOnce
+      expect(messageSpy).to.be.calledWithExactly({
+        comparison: comparisonPlotsFixture
+      })
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
+
     it('should open an image when receiving a plot zoomed message from the webview with a payload', async () => {
       const { webview, mockMessageReceived } = await buildPlotsWebview({
         disposer: disposable,

extension/src/webview/contract.ts

@@ -88,7 +88,8 @@ export enum MessageFromWebviewType {
   SHOW_LESS_COMMITS = 'show-less-commits',
   SWITCH_BRANCHES_VIEW = 'show-all-branches',
   SWITCH_COMMITS_VIEW = 'show-commits',
-  SELECT_BRANCHES = 'select-branches'
+  SELECT_BRANCHES = 'select-branches',
+  REFRESH_PLOTS = 'refresh-plots'
 }
 
 export type ColumnResizePayload = {
@@ -299,6 +300,7 @@ export type MessageFromWebview =
   | { type: MessageFromWebviewType.SWITCH_BRANCHES_VIEW }
   | { type: MessageFromWebviewType.SWITCH_COMMITS_VIEW }
   | { type: MessageFromWebviewType.SELECT_BRANCHES }
+  | { type: MessageFromWebviewType.REFRESH_PLOTS; payload: PlotsSection }
 
 export type MessageToWebview<T extends WebviewData> = {
   type: MessageToWebviewType.SET_DATA

webview/src/plots/components/App.test.tsx

@@ -45,8 +45,15 @@ import {
   TemplatePlotsById
 } from './plotDataStore'
 import { setMaxNbPlotsPerRow } from './webviewSlice'
-import { toggleDragAndDropMode as toggleTemplatePlotsDragAndDropMode } from './templatePlots/templatePlotsSlice'
-import { toggleDragAndDropMode as toggleCustomPlotsDragAndDropMode } from './customPlots/customPlotsSlice'
+import {
+  templatePlotsInitialState,
+  toggleDragAndDropMode as toggleTemplatePlotsDragAndDropMode
+} from './templatePlots/templatePlotsSlice'
+import {
+  customPlotsInitialState,
+  toggleDragAndDropMode as toggleCustomPlotsDragAndDropMode
+} from './customPlots/customPlotsSlice'
+import { comparisonTableInitialState } from './comparisonTable/comparisonTableSlice'
 import { plotsReducers, plotsStore } from '../store'
 import { vsCodeApi } from '../../shared/api'
 import {
@@ -700,42 +707,6 @@ describe('App', () => {
     expect(screen.queryByText('Images')).not.toBeInTheDocument()
   })
 
-  it('should hide plots when their section is collapsed (setting to null can break some Vega plots)', async () => {
-    renderAppWithOptionalData({
-      custom: customPlotsFixture
-    })
-
-    const summaryElement = await screen.findByText('Custom')
-    const visiblePlots = await screen.findAllByLabelText('Vega visualization')
-    for (const visiblePlot of visiblePlots) {
-      expect(visiblePlot).toBeInTheDocument()
-      expect(visiblePlot).toBeVisible()
-    }
-
-    fireEvent.click(summaryElement, {
-      bubbles: true,
-      cancelable: true
-    })
-
-    expect(mockPostMessage).toHaveBeenCalledWith({
-      payload: { [PlotsSection.CUSTOM_PLOTS]: true },
-      type: MessageFromWebviewType.TOGGLE_PLOTS_SECTION
-    })
-
-    sendSetDataMessage({
-      sectionCollapsed: {
-        ...DEFAULT_SECTION_COLLAPSED,
-        [PlotsSection.CUSTOM_PLOTS]: true
-      }
-    })
-
-    const hiddenPlots = await screen.findAllByLabelText('Vega visualization')
-    for (const hiddenPlot of hiddenPlots) {
-      expect(hiddenPlot).toBeInTheDocument()
-      expect(hiddenPlot).not.toBeVisible()
-    }
-  })
-
   it('should not toggle the custom plots section when its header is clicked and its title is selected', async () => {
     renderAppWithOptionalData({
       custom: customPlotsFixture
@@ -851,6 +822,95 @@ describe('App', () => {
     })
   })
 
+  it('should clear the template plots state when closing the section', () => {
+    const store = renderAppWithOptionalData({
+      template: templatePlotsFixture
+    })
+
+    const toggle = within(
+      screen.getByTestId(`${PlotsSection.TEMPLATE_PLOTS}-section-details`)
+    ).getAllByRole('button')[0]
+    fireEvent.click(toggle)
+
+    expect(store.getState().template).toMatchObject(templatePlotsInitialState)
+  })
+
+  it('should clear the custom plots state when closing the section', () => {
+    const store = renderAppWithOptionalData({
+      custom: customPlotsFixture
+    })
+
+    const toggle = within(
+      screen.getByTestId(`${PlotsSection.CUSTOM_PLOTS}-section-details`)
+    ).getAllByRole('button')[0]
+    fireEvent.click(toggle)
+
+    expect(store.getState().custom).toMatchObject(customPlotsInitialState)
+  })
+
+  it('should clear the comparison table state when closing the section', () => {
+    const store = renderAppWithOptionalData({
+      comparison: comparisonTableFixture
+    })
+
+    const toggle = within(
+      screen.getByTestId(`${PlotsSection.COMPARISON_TABLE}-section-details`)
+    ).getAllByRole('button')[0]
+    fireEvent.click(toggle)
+
+    expect(store.getState().comparison).toMatchObject(
+      comparisonTableInitialState
+    )
+  })
+
+  it('should send a message to refresh the template plots when closing the section', () => {
+    renderAppWithOptionalData({
+      template: templatePlotsFixture
+    })
+
+    const toggle = within(
+      screen.getByTestId(`${PlotsSection.TEMPLATE_PLOTS}-section-details`)
+    ).getAllByRole('button')[0]
+    fireEvent.click(toggle)
+
+    expect(mockPostMessage).toHaveBeenCalledWith({
+      payload: PlotsSection.TEMPLATE_PLOTS,
+      type: MessageFromWebviewType.REFRESH_PLOTS
+    })
+  })
+
+  it('should send a message to refresh the custom plots when closing the section', () => {
+    renderAppWithOptionalData({
+      custom: customPlotsFixture
+    })
+
+    const toggle = within(
+      screen.getByTestId(`${PlotsSection.CUSTOM_PLOTS}-section-details`)
+    ).getAllByRole('button')[0]
+    fireEvent.click(toggle)
+
+    expect(mockPostMessage).toHaveBeenCalledWith({
+      payload: PlotsSection.CUSTOM_PLOTS,
+      type: MessageFromWebviewType.REFRESH_PLOTS
+    })
+  })
+
+  it('should send a message to refresh the comparison table when closing the section', () => {
+    renderAppWithOptionalData({
+      comparison: comparisonTableFixture
+    })
+
+    const toggle = within(
+      screen.getByTestId(`${PlotsSection.COMPARISON_TABLE}-section-details`)
+    ).getAllByRole('button')[0]
+    fireEvent.click(toggle)
+
+    expect(mockPostMessage).toHaveBeenCalledWith({
+      payload: PlotsSection.COMPARISON_TABLE,
+      type: MessageFromWebviewType.REFRESH_PLOTS
+    })
+  })
+
   it('should display a slider to pick the number of items per row if there are items and the action is available', () => {
     const store = renderAppWithOptionalData({
       custom: customPlotsFixture

webview/src/plots/components/PlotsContainer.tsx

@@ -1,12 +1,12 @@
 import cx from 'classnames'
 import React, { useEffect, DetailedHTMLProps, HTMLAttributes } from 'react'
-import { useSelector } from 'react-redux'
+import { useDispatch, useSelector } from 'react-redux'
 import { PlotHeight, PlotsSection } from 'dvc/src/plots/webview/contract'
 import styles from './styles.module.scss'
 import { SizeSliders } from './SizeSliders'
-import { isDragAndDropModeSelector } from './util'
+import { clearStateActions, isDragAndDropModeSelector } from './util'
 import { PlotsState } from '../store'
-import { togglePlotsSection } from '../util/messages'
+import { refreshSection, togglePlotsSection } from '../util/messages'
 import { IconMenuItemProps } from '../../shared/components/iconMenu/IconMenuItem'
 import { Trash } from '../../shared/components/icons'
 import { SectionContainer } from '../../shared/components/sectionContainer/SectionContainer'
@@ -34,6 +34,7 @@ export const PlotsContainer: React.FC<PlotsContainerProps> = ({
   hasItems,
   noHeight
 }) => {
+  const dispatch = useDispatch()
   const open = !sectionCollapsed
   const maxNbPlotsPerRow = useSelector(
     (state: PlotsState) => state.webview.maxNbPlotsPerRow
@@ -57,7 +58,14 @@ export const PlotsContainer: React.FC<PlotsContainerProps> = ({
     })
   }
 
-  const toggleSection = () => togglePlotsSection(sectionKey, sectionCollapsed)
+  const toggleSection = (open: boolean) => {
+    togglePlotsSection(sectionKey, sectionCollapsed)
+
+    if (!open) {
+      dispatch(clearStateActions[sectionKey]())
+      refreshSection(sectionKey)
+    }
+  }
 
   return (
     <SectionContainer

webview/src/plots/components/comparisonTable/comparisonTableSlice.ts

@@ -47,6 +47,9 @@ export const comparisonTableSlice = createSlice({
     ) => {
       state.width = action.payload.nbItemsPerRowOrWidth
     },
+    clearState: () => {
+      return comparisonTableInitialState
+    },
     setCollapsed: (state, action: PayloadAction<boolean>) => {
       state.isCollapsed = action.payload
     },
@@ -72,7 +75,8 @@ export const {
   changeSize,
   changeDisabledDragIds,
   changeRowHeight,
-  toggleDragAndDropMode
+  toggleDragAndDropMode,
+  clearState
 } = comparisonTableSlice.actions
 
 export default comparisonTableSlice.reducer

webview/src/plots/components/customPlots/customPlotsSlice.ts

@@ -46,6 +46,10 @@ export const customPlotsSlice = createSlice({
       state.nbItemsPerRow = action.payload.nbItemsPerRowOrWidth
       state.height = action.payload.height
     },
+    clearState: () => {
+      removePlots([], PlotsSection.TEMPLATE_PLOTS)
+      return customPlotsInitialState
+    },
     setCollapsed: (state, action: PayloadAction<boolean>) => {
       state.isCollapsed = action.payload
     },
@@ -75,7 +79,12 @@ export const customPlotsSlice = createSlice({
   }
 })
 
-export const { update, setCollapsed, changeSize, toggleDragAndDropMode } =
-  customPlotsSlice.actions
+export const {
+  update,
+  setCollapsed,
+  changeSize,
+  toggleDragAndDropMode,
+  clearState
+} = customPlotsSlice.actions
 
 export default customPlotsSlice.reducer

webview/src/plots/components/templatePlots/templatePlotsSlice.ts

@@ -47,6 +47,10 @@ export const templatePlotsSlice = createSlice({
       state.nbItemsPerRow = action.payload.nbItemsPerRowOrWidth
       state.height = action.payload.height
     },
+    clearState: () => {
+      removePlots([], PlotsSection.TEMPLATE_PLOTS)
+      return templatePlotsInitialState
+    },
     setCollapsed: (state, action: PayloadAction<boolean>) => {
       state.isCollapsed = action.payload
     },
@@ -98,7 +102,8 @@ export const {
   setCollapsed,
   changeSize,
   toggleDragAndDropMode,
-  updateSections
+  updateSections,
+  clearState
 } = templatePlotsSlice.actions
 
 export default templatePlotsSlice.reducer