Fixes 6035: Add theme-aware lineage canvas controls (#32538)

* Fixes 6035: Add theme-aware lineage canvas controls

* test(ui): provide theme context in lineage tests

Author: shah-harshit

openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/CanvasEdgeRenderer.component.test.tsx

@@ -10,8 +10,9 @@
  *  See the License for the specific language governing permissions and
  *  limitations under the License.
  */
-import { fireEvent, render, waitFor } from '@testing-library/react';
+import { act, fireEvent, render, waitFor } from '@testing-library/react';
 import { Edge } from 'reactflow';
+import { ThemeProvider } from '../../../context/UntitledUIThemeProvider/theme-provider';
 import { CanvasEdgeRenderer } from './CanvasEdgeRenderer.component';
 
 const mockRedraw = jest.fn();
@@ -130,7 +131,7 @@ describe('CanvasEdgeRenderer', () => {
     reactFlowContainer.appendChild(wrapper);
     document.body.appendChild(reactFlowContainer);
 
-    return render(ui, { container: wrapper });
+    return render(ui, { container: wrapper, wrapper: ThemeProvider });
   };
 
   it('renders canvas element', () => {
@@ -459,7 +460,7 @@ describe('CanvasEdgeRenderer', () => {
         },
       ] as unknown as ResizeObserverEntry[];
 
-      resizeCallback(mockEntries, {} as ResizeObserver);
+      act(() => resizeCallback?.(mockEntries, {} as ResizeObserver));
     }
 
     await waitFor(() => {

openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/CanvasEdgeRenderer.component.tsx

@@ -14,6 +14,7 @@ import React, { useEffect, useMemo, useRef, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { Edge, useReactFlow, useViewport } from 'reactflow';
 import { useLineageProvider } from '../../../context/LineageProvider/LineageProvider';
+import { useTheme } from '../../../context/UntitledUIThemeProvider/theme-provider';
 import { useCanvasEdgeRenderer } from '../../../hooks/useCanvasEdgeRenderer';
 import { useCanvasMouseEvents } from '../../../hooks/useCanvasMouseEvents';
 import { useLineageEdgeColors } from '../../../hooks/useLineageEdgeColors';
@@ -39,6 +40,7 @@ export const CanvasEdgeRenderer: React.FC<CanvasEdgeRendererProps> = ({
   hoverEdge,
 }) => {
   const { t } = useTranslation();
+  const { theme } = useTheme();
   const edgeColors = useLineageEdgeColors();
   const canvasRef = useRef<HTMLCanvasElement>(null);
   const containerRef = useRef<HTMLDivElement>(null);
@@ -114,6 +116,7 @@ export const CanvasEdgeRenderer: React.FC<CanvasEdgeRendererProps> = ({
     hoverEdge,
     containerWidth: containerSize.width,
     containerHeight: containerSize.height,
+    theme,
   });
 
   useEffect(() => {

openmetadata-ui/src/main/resources/ui/src/components/Lineage/Lineage.test.tsx

@@ -10,10 +10,15 @@
  *  See the License for the specific language governing permissions and
  *  limitations under the License.
  */
-import { fireEvent, render, screen } from '@testing-library/react';
+import {
+  fireEvent,
+  render as renderComponent,
+  screen,
+} from '@testing-library/react';
 import { DragEvent } from 'react';
 import ReactFlow from 'reactflow';
 import { useLineageProvider } from '../../context/LineageProvider/LineageProvider';
+import { ThemeProvider } from '../../context/UntitledUIThemeProvider/theme-provider';
 import { EntityType } from '../../enums/entity.enum';
 import { Table } from '../../generated/entity/data/table';
 import { useLineageStore } from '../../hooks/useLineageStore';
@@ -150,6 +155,9 @@ const mockLineageStore = {
   activeLayer: [],
 };
 
+const render = (component: React.ReactElement) =>
+  renderComponent(component, { wrapper: ThemeProvider });
+
 jest.mock('../../context/LineageProvider/LineageProvider', () => ({
   useLineageProvider: jest.fn(),
 }));

openmetadata-ui/src/main/resources/ui/src/hooks/useCanvasEdgeRenderer.test.ts

@@ -167,6 +167,7 @@ describe('useCanvasEdgeRenderer', () => {
         colors: createMockColors(),
         containerWidth: 800,
         containerHeight: 600,
+        theme: 'light',
       })
     );
 
@@ -183,6 +184,7 @@ describe('useCanvasEdgeRenderer', () => {
         colors: createMockColors(),
         containerWidth: 800,
         containerHeight: 600,
+        theme: 'light',
       })
     );
 
@@ -191,6 +193,32 @@ describe('useCanvasEdgeRenderer', () => {
     expect(requestAnimationFrame).toHaveBeenCalled();
   });
 
+  it('schedules a redraw when the active theme changes', () => {
+    const colors = createMockColors();
+    const dqHighlightedEdges = new Set<string>();
+    const edges: Edge[] = [];
+    const initialProps: { theme: 'light' | 'dark' } = { theme: 'light' };
+    const { rerender } = renderHook(
+      ({ theme }: { theme: 'light' | 'dark' }) =>
+        useCanvasEdgeRenderer({
+          canvasRef,
+          colors,
+          containerHeight: 600,
+          containerWidth: 800,
+          dqHighlightedEdges,
+          edges,
+          theme,
+        }),
+      { initialProps }
+    );
+    const initialRedrawCount = (requestAnimationFrame as jest.Mock).mock.calls
+      .length;
+
+    rerender({ theme: 'dark' });
+
+    expect(requestAnimationFrame).toHaveBeenCalledTimes(initialRedrawCount + 1);
+  });
+
   it('draws visible edges', () => {
     const edge = createMockEdge();
     const node1 = createMockNode('node-1');
@@ -215,6 +243,7 @@ describe('useCanvasEdgeRenderer', () => {
         colors: createMockColors(),
         containerWidth: 800,
         containerHeight: 600,
+        theme: 'light',
       })
     );
 
@@ -235,6 +264,7 @@ describe('useCanvasEdgeRenderer', () => {
         colors: createMockColors(),
         containerWidth: 800,
         containerHeight: 600,
+        theme: 'light',
       })
     );
 
@@ -263,6 +293,7 @@ describe('useCanvasEdgeRenderer', () => {
         colors: createMockColors(),
         containerWidth: 800,
         containerHeight: 600,
+        theme: 'light',
       })
     );
 
@@ -305,6 +336,7 @@ describe('useCanvasEdgeRenderer', () => {
         colors: createMockColors(),
         containerWidth: 800,
         containerHeight: 600,
+        theme: 'light',
       })
     );
 
@@ -343,6 +375,7 @@ describe('useCanvasEdgeRenderer', () => {
         colors: createMockColors(),
         containerWidth: 800,
         containerHeight: 600,
+        theme: 'light',
       })
     );
 
@@ -392,6 +425,7 @@ describe('useCanvasEdgeRenderer', () => {
         colors: createMockColors(),
         containerWidth: 800,
         containerHeight: 600,
+        theme: 'light',
       })
     );
 
@@ -447,6 +481,7 @@ describe('useCanvasEdgeRenderer', () => {
         colors: createMockColors(),
         containerWidth: 800,
         containerHeight: 600,
+        theme: 'light',
       })
     );
 
@@ -465,6 +500,7 @@ describe('useCanvasEdgeRenderer', () => {
         colors: createMockColors(),
         containerWidth: 800,
         containerHeight: 600,
+        theme: 'light',
       })
     );
 
@@ -502,6 +538,7 @@ describe('useCanvasEdgeRenderer', () => {
         colors: createMockColors(),
         containerWidth: 800,
         containerHeight: 600,
+        theme: 'light',
       })
     );
 

openmetadata-ui/src/main/resources/ui/src/hooks/useCanvasEdgeRenderer.ts

@@ -20,12 +20,15 @@ import {
 } from 'react';
 import type { Edge } from 'reactflow';
 import { Position, useNodes, useReactFlow, useViewport } from 'reactflow';
+import { Theme } from '../context/UntitledUIThemeProvider/theme-provider.interface';
 import {
   CanvasButton,
+  CanvasButtonColors,
   createCanvasButton,
   drawCanvasButton,
   ECanvasButtonType,
   isPointInButton,
+  resolveCanvasButtonColors,
 } from '../utils/CanvasButtonUtils';
 import {
   drawArrowMarker,
@@ -51,6 +54,7 @@ interface UseCanvasEdgeRendererProps {
   colors: LineageEdgeColors;
   containerWidth: number;
   containerHeight: number;
+  theme: Theme;
 }
 
 interface EdgeHitEntry {
@@ -99,7 +103,8 @@ const getCanvasButtonHit = (
     >
   >,
   hoveredButtonRef: MutableRefObject<CanvasButton | null>,
-  isDQEnabled: boolean
+  isDQEnabled: boolean,
+  getButtonColors: () => CanvasButtonColors
 ): CanvasButtonHitData | null => {
   const edgeData = edge.data ?? {};
   const { hasPipeline, hasFunction } = getEdgeButtonFlags(edgeData);
@@ -128,7 +133,13 @@ const getCanvasButtonHit = (
     hoveredButtonRef.current?.type === button.type;
 
   ctx.save();
-  drawCanvasButton(ctx, button, isButtonHovered, isDQEnabled);
+  drawCanvasButton(
+    ctx,
+    button,
+    getButtonColors(),
+    isButtonHovered,
+    isDQEnabled
+  );
   ctx.restore();
 
   return { button, edge };
@@ -142,6 +153,7 @@ export function useCanvasEdgeRenderer({
   colors,
   containerWidth,
   containerHeight,
+  theme,
 }: UseCanvasEdgeRendererProps) {
   const rafIdRef = useRef<number>();
   const isDirtyRef = useRef(false);
@@ -353,6 +365,14 @@ export function useCanvasEdgeRenderer({
 
     const hitPaths: EdgeHitEntry[] = [];
     const canvasButtons: CanvasButtonHitData[] = [];
+    // Resolve lazily inside the scheduled draw so the root theme class is current,
+    // while graphs without pipeline/function controls avoid unnecessary DOM probes.
+    let buttonColors: CanvasButtonColors | undefined;
+    const getButtonColors = () => {
+      buttonColors ??= resolveCanvasButtonColors();
+
+      return buttonColors;
+    };
 
     visibleEdges.forEach((edge) => {
       ctx.save();
@@ -368,7 +388,8 @@ export function useCanvasEdgeRenderer({
         edge,
         edgePathCacheRef,
         hoveredButtonRef,
-        isDQEnabled
+        isDQEnabled,
+        getButtonColors
       );
 
       if (buttonHit) {
@@ -507,6 +528,7 @@ export function useCanvasEdgeRenderer({
     selectedColumn,
     dqHighlightedEdges,
     colors,
+    theme,
   ]);
 
   useEffect(() => {