[Traces][Discover] Prevent flyout remount when switching document types in Trace Waterfall

src/platform/plugins/shared/unified_doc_viewer/public/components/observability/traces/components/full_screen_waterfall/index.test.tsx

@@ -15,22 +15,25 @@ import { setUnifiedDocViewerServices } from '../../../../../plugin';
 import type { UnifiedDocViewerServices } from '../../../../../types';
 
 jest.mock('./waterfall_flyout/span_flyout', () => ({
-  SpanFlyout: ({ traceId, spanId, _, activeSection }: any) => (
-    <div
-      data-test-subj="spanFlyout"
-      data-trace-id={traceId}
-      data-span-id={spanId}
-      data-active-section={activeSection}
-    />
-  ),
   spanFlyoutId: 'spanFlyout',
 }));
 
 jest.mock('./waterfall_flyout/logs_flyout', () => ({
-  LogsFlyout: ({ id, _ }: any) => <div data-test-subj="logsFlyout" data-id={id} />,
   logsFlyoutId: 'logsFlyout',
 }));
 
+jest.mock('./waterfall_flyout/document_detail_flyout', () => ({
+  DocumentDetailFlyout: ({ type, docId, traceId, activeSection }: any) => (
+    <div
+      data-test-subj={type === 'spanFlyout' ? 'spanFlyout' : 'logsFlyout'}
+      data-trace-id={traceId}
+      data-span-id={docId}
+      data-id={docId}
+      data-active-section={activeSection}
+    />
+  ),
+}));
+
 describe('FullScreenWaterfall', () => {
   const defaultProps: FullScreenWaterfallProps = {
     traceId: 'test-trace-id',

src/platform/plugins/shared/unified_doc_viewer/public/components/observability/traces/components/full_screen_waterfall/index.tsx

@@ -20,10 +20,9 @@ import type { DocViewRenderProps } from '@kbn/unified-doc-viewer/types';
 import React, { useCallback, useState } from 'react';
 import { getUnifiedDocViewerServices } from '../../../../../plugin';
 import type { TraceOverviewSections } from '../../doc_viewer_overview/overview';
-import type { logsFlyoutId as logsFlyoutIdType } from './waterfall_flyout/logs_flyout';
-import { LogsFlyout, logsFlyoutId } from './waterfall_flyout/logs_flyout';
-import type { spanFlyoutId as spanFlyoutIdType } from './waterfall_flyout/span_flyout';
-import { SpanFlyout, spanFlyoutId } from './waterfall_flyout/span_flyout';
+import { spanFlyoutId } from './waterfall_flyout/span_flyout';
+import { logsFlyoutId } from './waterfall_flyout/logs_flyout';
+import { DocumentDetailFlyout, type DocumentType } from './waterfall_flyout/document_detail_flyout';
 
 export const EUI_FLYOUT_BODY_OVERFLOW_CLASS = 'euiFlyoutBody__overflow';
 
@@ -51,9 +50,7 @@ export const FullScreenWaterfall = ({
   const { euiTheme } = useEuiTheme();
   const [docId, setDocId] = useState<string | null>(null);
   const [docIndex, setDocIndex] = useState<string | undefined>(undefined);
-  const [activeFlyoutId, setActiveFlyoutId] = useState<
-    typeof spanFlyoutIdType | typeof logsFlyoutIdType | null
-  >(null);
+  const [activeFlyoutType, setActiveFlyoutType] = useState<DocumentType | null>(null);
   const [activeSection, setActiveSection] = useState<TraceOverviewSections | undefined>();
   const [scrollElement, setScrollElement] = useState<Element | null>(null);
 
@@ -88,7 +85,7 @@ export const FullScreenWaterfall = ({
   }, []);
 
   function handleCloseFlyout() {
-    setActiveFlyoutId(null);
+    setActiveFlyoutType(null);
     setActiveSection(undefined);
     setDocId(null);
     setDocIndex(undefined);
@@ -98,7 +95,7 @@ export const FullScreenWaterfall = ({
     setActiveSection(undefined);
     setDocId(nodeSpanId);
     setDocIndex(undefined);
-    setActiveFlyoutId(spanFlyoutId);
+    setActiveFlyoutType(spanFlyoutId);
   }
 
   function handleErrorClick(params: {
@@ -109,12 +106,12 @@ export const FullScreenWaterfall = ({
     docIndex?: string;
   }) {
     if (params.errorCount > 1) {
-      setActiveFlyoutId(spanFlyoutId);
+      setActiveFlyoutType(spanFlyoutId);
       setActiveSection('errors-table');
       setDocId(params.docId);
       setDocIndex(undefined);
     } else if (params.errorDocId) {
-      setActiveFlyoutId(logsFlyoutId);
+      setActiveFlyoutType(logsFlyoutId);
       setDocId(params.errorDocId);
       setDocIndex(params.docIndex);
     }
@@ -163,23 +160,16 @@ export const FullScreenWaterfall = ({
         </div>
       </EuiFlyoutBody>
 
-      {docId && activeFlyoutId ? (
-        activeFlyoutId === spanFlyoutId ? (
-          <SpanFlyout
-            traceId={traceId}
-            spanId={docId}
-            dataView={dataView}
-            onCloseFlyout={handleCloseFlyout}
-            activeSection={activeSection}
-          />
-        ) : (
-          <LogsFlyout
-            onCloseFlyout={handleCloseFlyout}
-            id={docId}
-            index={docIndex}
-            dataView={dataView}
-          />
-        )
+      {docId && activeFlyoutType ? (
+        <DocumentDetailFlyout
+          type={activeFlyoutType}
+          docId={docId}
+          docIndex={docIndex}
+          traceId={traceId}
+          dataView={dataView}
+          onCloseFlyout={handleCloseFlyout}
+          activeSection={activeSection}
+        />
       ) : null}
     </EuiFlyout>
   );