test(app-routes): add tests for private route

frontend/src/AppRoutes/Private.tsx

@@ -101,6 +101,22 @@ function PrivateRoute({ children }: PrivateRouteProps): JSX.Element {
 					preference.name === ORG_PREFERENCES.ORG_ONBOARDING,
 			)?.value;
 
+			// Don't redirect to onboarding if workspace has issues (blocked, suspended, or restricted)
+			// User needs access to settings/billing to fix payment issues
+			const isWorkspaceBlocked = trialInfo?.workSpaceBlock;
+			const isWorkspaceSuspended = activeLicense?.state === LicenseState.DEFAULTED;
+			const isWorkspaceAccessRestricted =
+				activeLicense?.state === LicenseState.TERMINATED ||
+				activeLicense?.state === LicenseState.EXPIRED ||
+				activeLicense?.state === LicenseState.CANCELLED;
+
+			const hasWorkspaceIssue =
+				isWorkspaceBlocked || isWorkspaceSuspended || isWorkspaceAccessRestricted;
+
+			if (hasWorkspaceIssue) {
+				return;
+			}
+
 			const isFirstUser = checkFirstTimeUser();
 			if (
 				isFirstUser &&
@@ -119,40 +135,36 @@ function PrivateRoute({ children }: PrivateRouteProps): JSX.Element {
 		orgPreferences,
 		usersData,
 		pathname,
+		trialInfo?.workSpaceBlock,
+		activeLicense?.state,
 	]);
 
-	const navigateToWorkSpaceBlocked = (route: any): void => {
-		const { path } = route;
-
+	const navigateToWorkSpaceBlocked = useCallback((): void => {
 		const isRouteEnabledForWorkspaceBlockedState =
 			isAdmin &&
-			(path === ROUTES.SETTINGS ||
-				path === ROUTES.ORG_SETTINGS ||
-				path === ROUTES.MEMBERS_SETTINGS ||
-				path === ROUTES.BILLING ||
-				path === ROUTES.MY_SETTINGS);
+			(pathname === ROUTES.SETTINGS ||
+				pathname === ROUTES.ORG_SETTINGS ||
+				pathname === ROUTES.MEMBERS_SETTINGS ||
+				pathname === ROUTES.BILLING ||
+				pathname === ROUTES.MY_SETTINGS);
 
 		if (
-			path &&
-			path !== ROUTES.WORKSPACE_LOCKED &&
+			pathname &&
+			pathname !== ROUTES.WORKSPACE_LOCKED &&
 			!isRouteEnabledForWorkspaceBlockedState
 		) {
 			history.push(ROUTES.WORKSPACE_LOCKED);
 		}
-	};
+	}, [isAdmin, pathname]);
 
-	const navigateToWorkSpaceAccessRestricted = (route: any): void => {
-		const { path } = route;
-
-		if (path && path !== ROUTES.WORKSPACE_ACCESS_RESTRICTED) {
+	const navigateToWorkSpaceAccessRestricted = useCallback((): void => {
+		if (pathname && pathname !== ROUTES.WORKSPACE_ACCESS_RESTRICTED) {
 			history.push(ROUTES.WORKSPACE_ACCESS_RESTRICTED);
 		}
-	};
+	}, [pathname]);
 
 	useEffect(() => {
 		if (!isFetchingActiveLicense && activeLicense) {
-			const currentRoute = mapRoutes.get('current');
-
 			const isTerminated = activeLicense.state === LicenseState.TERMINATED;
 			const isExpired = activeLicense.state === LicenseState.EXPIRED;
 			const isCancelled = activeLicense.state === LicenseState.CANCELLED;
@@ -161,61 +173,53 @@ function PrivateRoute({ children }: PrivateRouteProps): JSX.Element {
 
 			const { platform } = activeLicense;
 
-			if (
-				isWorkspaceAccessRestricted &&
-				platform === LicensePlatform.CLOUD &&
-				currentRoute
-			) {
-				navigateToWorkSpaceAccessRestricted(currentRoute);
+			if (isWorkspaceAccessRestricted && platform === LicensePlatform.CLOUD) {
+				navigateToWorkSpaceAccessRestricted();
 			}
 		}
-	}, [isFetchingActiveLicense, activeLicense, mapRoutes, pathname]);
+	}, [
+		isFetchingActiveLicense,
+		activeLicense,
+		navigateToWorkSpaceAccessRestricted,
+	]);
 
 	useEffect(() => {
 		if (!isFetchingActiveLicense) {
-			const currentRoute = mapRoutes.get('current');
 			const shouldBlockWorkspace = trialInfo?.workSpaceBlock;
 
 			if (
 				shouldBlockWorkspace &&
-				currentRoute &&
 				activeLicense?.platform === LicensePlatform.CLOUD
 			) {
-				navigateToWorkSpaceBlocked(currentRoute);
+				navigateToWorkSpaceBlocked();
 			}
 		}
-		// eslint-disable-next-line react-hooks/exhaustive-deps
 	}, [
 		isFetchingActiveLicense,
 		trialInfo?.workSpaceBlock,
 		activeLicense?.platform,
-		mapRoutes,
-		pathname,
+		navigateToWorkSpaceBlocked,
 	]);
 
-	const navigateToWorkSpaceSuspended = (route: any): void => {
-		const { path } = route;
-
-		if (path && path !== ROUTES.WORKSPACE_SUSPENDED) {
+	const navigateToWorkSpaceSuspended = useCallback((): void => {
+		if (pathname && pathname !== ROUTES.WORKSPACE_SUSPENDED) {
 			history.push(ROUTES.WORKSPACE_SUSPENDED);
 		}
-	};
+	}, [pathname]);
 
 	useEffect(() => {
 		if (!isFetchingActiveLicense && activeLicense) {
-			const currentRoute = mapRoutes.get('current');
 			const shouldSuspendWorkspace =
 				activeLicense.state === LicenseState.DEFAULTED;
 
 			if (
 				shouldSuspendWorkspace &&
-				currentRoute &&
 				activeLicense.platform === LicensePlatform.CLOUD
 			) {
-				navigateToWorkSpaceSuspended(currentRoute);
+				navigateToWorkSpaceSuspended();
 			}
 		}
-	}, [isFetchingActiveLicense, activeLicense, mapRoutes, pathname]);
+	}, [isFetchingActiveLicense, activeLicense, navigateToWorkSpaceSuspended]);
 
 	useEffect(() => {
 		if (org && org.length > 0 && org[0].id !== undefined) {

pkg/query-service/app/traces/tracedetail/waterfall.go

@@ -1,6 +1,7 @@
 package tracedetail
 
 import (
+	"maps"
 	"slices"
 	"sort"
 
@@ -63,26 +64,22 @@ func findIndexForSelectedSpanFromPreOrder(spans []*model.Span, selectedSpanId st
 	return selectedSpanIndex
 }
 
-func getPathFromRootToSelectedSpanId(node *model.Span, selectedSpanId string, uncollapsedSpans []string, isSelectedSpanIDUnCollapsed bool) (bool, []string) {
+func getPathFromRootToSelectedSpanId(node *model.Span, selectedSpanId string) (bool, []string) {
 	spansFromRootToNode := []string{}
 
+	spansFromRootToNode = append(spansFromRootToNode, node.SpanID)
 	if node.SpanID == selectedSpanId {
-		if isSelectedSpanIDUnCollapsed && !slices.Contains(uncollapsedSpans, node.SpanID) {
-			spansFromRootToNode = append(spansFromRootToNode, node.SpanID)
-		}
 		return true, spansFromRootToNode
 	}
 
 	isPresentInSubtreeForTheNode := false
 	for _, child := range node.Children {
-		isPresentInThisSubtree, _spansFromRootToNode := getPathFromRootToSelectedSpanId(child, selectedSpanId, uncollapsedSpans, isSelectedSpanIDUnCollapsed)
+		isPresentInThisSubtree, _spansFromRootToNode := getPathFromRootToSelectedSpanId(child, selectedSpanId)
 		// if the interested node is present in the given subtree then add the span node to uncollapsed node list
 		if isPresentInThisSubtree {
-			if !slices.Contains(uncollapsedSpans, node.SpanID) {
-				spansFromRootToNode = append(spansFromRootToNode, node.SpanID)
-			}
 			isPresentInSubtreeForTheNode = true
 			spansFromRootToNode = append(spansFromRootToNode, _spansFromRootToNode...)
+			break
 		}
 	}
 	return isPresentInSubtreeForTheNode, spansFromRootToNode
@@ -92,7 +89,7 @@ func getPathFromRootToSelectedSpanId(node *model.Span, selectedSpanId string, un
 // throughout the recursion. Per-call state (level, isPartOfPreOrder, etc.)
 // is passed as direct arguments.
 type traverseOpts struct {
-	uncollapsedSpans []string
+	uncollapsedSpans map[string]struct{}
 	selectedSpanID   string
 }
 
@@ -134,7 +131,7 @@ func traverseTrace(span *model.Span, opts traverseOpts, level uint64, isPartOfPr
 		preOrderTraversal = append(preOrderTraversal, &nodeWithoutChildren)
 	}
 
-	isAlreadyUncollapsed := slices.Contains(opts.uncollapsedSpans, span.SpanID)
+	_, isAlreadyUncollapsed := opts.uncollapsedSpans[span.SpanID]
 	for index, child := range span.Children {
 		_childTraversal := traverseTrace(child, opts, level+1, isPartOfPreOrder && isAlreadyUncollapsed, index != (len(span.Children)-1))
 		preOrderTraversal = append(preOrderTraversal, _childTraversal...)
@@ -169,16 +166,28 @@ func GetSelectedSpans(uncollapsedSpans []string, selectedSpanID string, traceRoo
 
 	var preOrderTraversal = make([]*model.Span, 0)
 	var rootServiceName, rootServiceEntryPoint string
-	updatedUncollapsedSpans := uncollapsedSpans
+
+	// create a map of uncollapsed spans for quick lookup
+	uncollapsedSpanMap := make(map[string]struct{})
+	for _, spanID := range uncollapsedSpans {
+		uncollapsedSpanMap[spanID] = struct{}{}
+	}
 
 	selectedSpanIndex := -1
 	for _, rootSpanID := range traceRoots {
 		if rootNode, exists := spanIdToSpanNodeMap[rootSpanID.SpanID]; exists {
-			_, spansFromRootToNode := getPathFromRootToSelectedSpanId(rootNode, selectedSpanID, updatedUncollapsedSpans, isSelectedSpanIDUnCollapsed)
-			updatedUncollapsedSpans = append(updatedUncollapsedSpans, spansFromRootToNode...)
+			present, spansFromRootToNode := getPathFromRootToSelectedSpanId(rootNode, selectedSpanID)
+			if present {
+				for _, spanID := range spansFromRootToNode {
+					if selectedSpanID == spanID && !isSelectedSpanIDUnCollapsed {
+						continue
+					}
+					uncollapsedSpanMap[spanID] = struct{}{}
+				}
+			}
 
 			opts := traverseOpts{
-				uncollapsedSpans: updatedUncollapsedSpans,
+				uncollapsedSpans: uncollapsedSpanMap,
 				selectedSpanID:   selectedSpanID,
 			}
 			_preOrderTraversal := traverseTrace(rootNode, opts, 0, true, false)
@@ -223,5 +232,5 @@ func GetSelectedSpans(uncollapsedSpans []string, selectedSpanID string, traceRoo
 		startIndex = 0
 	}
 
-	return preOrderTraversal[startIndex:endIndex], updatedUncollapsedSpans, rootServiceName, rootServiceEntryPoint
+	return preOrderTraversal[startIndex:endIndex], slices.Collect(maps.Keys(uncollapsedSpanMap)), rootServiceName, rootServiceEntryPoint
 }