Fix Journey Canvas tests on top of cl/705410457

src/app/demo_page/demo_page.ng.html

@@ -556,6 +556,30 @@
         </select>
       </mat-form-field>
     </div>
+    <mat-divider />
+    <div class="demo-control-group">
+      <p class="demo-control-header">
+        Node selection
+      </p>
+      <mat-form-field
+        appearance="fill"
+        class="demo-select-full-width"
+      >
+        <mat-label>Node/Group</mat-label>
+        <mat-select
+          [(ngModel)]="selectedNode"
+          [compareWith]="selectedNodeComparator"
+          id="node-select"
+        >
+          <mat-option
+            *ngFor="let selectable of allSelectableNodesAndGroups"
+            [value]="selectable"
+          >
+            {{selectable.label}}
+          </mat-option>
+        </mat-select>
+      </mat-form-field>
+    </div>
   </div>
 </div>
 

src/app/demo_page/demo_page.scss

@@ -71,6 +71,9 @@ h2 {
       &.demo-select {
         width: 40%;
       }
+      &.demo-select-full-width {
+        width: 80%;
+      }
     }
 
     .demo-control-group {

src/app/demo_page/demo_page.spec.ts

@@ -113,6 +113,21 @@ describe('Demo Page', () => {
       await screenShot.expectMatch('dark-mode');
     });
   });
+
+  describe('Node selection', () => {
+    it('Selects correct nested iteration loop node (screenshot)', async () => {
+      await harness.getNodeSelectInput().then(
+          select => select.clickOptions(
+              {text: 'Trainer (TensorFlow Training, it-1)'}));
+      await screenShot.expectMatch('select-nested-iteration-loop');
+    });
+
+    it('Selects correct nested group node (screenshot)', async () => {
+      await harness.getNodeSelectInput().then(
+          select => select.clickOptions({text: 'Fake Exec 1 (sub1)'}));
+      await screenShot.expectMatch('select-nested-group');
+    })
+  })
 });
 
 

src/app/demo_page/demo_page.ts

@@ -20,6 +20,7 @@ import {HttpClientModule} from '@angular/common/http';
 import {Component, ElementRef, NgModule, TemplateRef, ViewChild, ViewEncapsulation} from '@angular/core';
 import {FormsModule, ReactiveFormsModule} from '@angular/forms';
 import {MatButtonModule} from '@angular/material/button';
+import {MatOptionModule} from '@angular/material/core';
 import {MatDialog, MatDialogModule} from '@angular/material/dialog';
 import {MatDividerModule} from '@angular/material/divider';
 import {MatFormFieldModule} from '@angular/material/form-field';
@@ -170,6 +171,32 @@ function translateColor(col: string, isBg = false) {
   }
 }
 
+function getAllSelectableNestedNodesAndGroups(
+    params: {nodes: DagNode[], groups: DagGroup[]},
+    path: string[] = []): Array<SelectedNode&{label: string}> {
+  const allNodesAndGroups = [];
+  const pathLabel = path.length ? ` (${path.join(', ')})` : '';
+  for (const subGroup of params.groups) {
+    allNodesAndGroups.push(
+        {node: subGroup, path, label: `${subGroup.id}${pathLabel}`});
+    if (!subGroup.treatAsLoop) {
+      allNodesAndGroups.push(...getAllSelectableNestedNodesAndGroups(
+          subGroup, [...path, subGroup.id]));
+    } else {
+      // We need to skip the iteration node itself, as it is not selectable
+      subGroup.groups.forEach((iterationGroup) => {
+        allNodesAndGroups.push(...getAllSelectableNestedNodesAndGroups(
+            iterationGroup, [...path, subGroup.id, iterationGroup.id]));
+      });
+    }
+  }
+  for (const subNode of params.nodes) {
+    allNodesAndGroups.push(
+        {node: subNode, path, label: `${subNode.id}${pathLabel}`});
+  }
+  return allNodesAndGroups;
+}
+
 /** Demo component for directed acyclic graph view. */
 @Component({
   standalone: false,
@@ -237,6 +264,7 @@ export class DagDemoPage {
   userConfigChange = new Subject<UserConfig>();
   destroy = new Subject<void>();
 
+  allSelectableNodesAndGroups: Array<SelectedNode&{label: string}> = [];
   constructor(public dialog: MatDialog) {
     this.setCurrDataset(DEFAULT_DATASET, true);
     this.resetAll();
@@ -339,6 +367,8 @@ export class DagDemoPage {
     const newGraph = cloneGraph(this.datasets[name]);
     this.calibrateNodes(newGraph);
     this.currDataset = newGraph;
+    this.allSelectableNodesAndGroups =
+        getAllSelectableNestedNodesAndGroups(newGraph);
   }
 
   onDatasetChanged(event: Event) {
@@ -531,6 +561,14 @@ export class DagDemoPage {
     this.destroy.next();
     this.destroy.complete();
   }
+
+  triggerSelection(node: SelectedNode) {
+    this.selectedNode = node;
+  }
+
+  selectedNodeComparator(a: SelectedNode, b: SelectedNode) {
+    return a.node.id === b.node.id && a.path.join('') === b.path.join('');
+  }
 }
 
 @NgModule({
@@ -546,6 +584,7 @@ export class DagDemoPage {
     DagScaffoldModule,
     DagToolbarModule,
     MatSelectModule,
+    MatOptionModule,
     WorkflowGraphIconModule,
     MatDialogModule,
     FormsModule,

src/app/directed_acyclic_graph.ts

@@ -29,7 +29,7 @@ import {ColorThemeLoader} from './color_theme_loader';
 import {DagStateService} from './dag-state.service';
 import {STATE_SERVICE_PROVIDER} from './dag-state.service.provider';
 import {baseColors, BLUE_THEME, clampVal, CLASSIC_THEME, createDAGFeatures, createNewSizeConfig, type DagTheme, DEFAULT_LAYOUT_OPTIONS, DEFAULT_THEME, defaultFeatures, defaultZoomConfig, EdgeStyle, type FeatureToggleOptions, generateTheme, getMargin, isPoint, type LayoutOptions, type Logger, MarkerStyle, type MinimapPosition, nanSafePt, NODE_HEIGHT, NODE_WIDTH, NodeState, OrientationMarginConfig, RankAlignment, RankDirection, RankerAlgorithim, SCROLL_STEP_PER_DELTA, SizeConfig, SVG_ELEMENT_SIZE, type ZoomConfig} from './data_types_internal';
-import {DagRaw, DagRawModule, EnhancedDagGroup, GraphDims} from './directed_acyclic_graph_raw';
+import {DagRaw, DagRawModule, EnhancedDagGroup, GraphDims, setEnhancedGroupSelection} from './directed_acyclic_graph_raw';
 import {DagLogger} from './logger/dag_logger';
 import {Minimap, MinimapModule} from './minimap/minimap';
 import {type DagEdge, DagGroup, DagNode, GraphSpec, GroupIterationRecord, GroupToggleEvent, isDagreInit, NodeMap, type NodeRef, Point, type SelectedNode} from './node_spec';
@@ -577,7 +577,9 @@ export class DirectedAcyclicGraph implements OnInit, OnDestroy {
           const loopGroup = parent?.$groups?.find(g => g.id === segment);
           if (loopGroup) {
             const selectedLoop = path.shift()!;
-            loopGroup.selectedLoopId = selectedLoop;
+            const selectedLoopNode =
+                loopGroup.groups.find(n => n.id === selectedLoop);
+            setEnhancedGroupSelection(loopGroup, selectedLoopNode!);
             // Increase pathDepth as we've consumed 2 path segments
             pathDepth++;
             return true;
@@ -592,6 +594,8 @@ export class DirectedAcyclicGraph implements OnInit, OnDestroy {
         expandedStateChanges++;
         // Detect changes so groups can expand properly
         this.detectChanges();
+        // Detect changes in new dagEl so QueryList is updated properly
+        dagEl.detectChanges();
       }
       pathDepth++;
     }

src/app/directed_acyclic_graph_raw.ng.html

@@ -1,5 +1,5 @@
 <aside
-  *ngIf="!noEmptySpaceAlloc"
+  *ngIf="!noEmptySpaceAlloc && visible"
   aria-hidden="true"
   class="empty-space-alloc"
   [style.width.px]="graphWidth"
@@ -9,6 +9,7 @@
   [ngClass]="['dag-component', toggleClass(collapsed, 'collapse')]"
   [style.width.px]="graphWidth"
   [style.height.px]="graphHeight"
+  *ngIf="visible"
 >
   <svg class="edge canvas" aria-hidden="true">
     <ng-container *ngFor="let e of edges; trackBy: edgeTrack">
@@ -351,6 +352,7 @@
           [groups]="group.groups"
           (graphResize)="storeSubDagDims($event, group)"
           [resolveReference]="resolveReference"
+          [visible]="isGroupExpanded(group)"
         ></ai-dag-raw>
         <ng-container *ngIf="!!(group.treatAsLoop && group._cachedSelection)">
           <ai-dag-raw
@@ -362,6 +364,7 @@
             [nodes]="[group._cachedSelection]"
             (graphResize)="storeSubDagDims($event, group)"
             [resolveReference]="resolveReference"
+            [visible]="isGroupExpanded(group)"
           ></ai-dag-raw>
           <ai-dag-raw
             #subDag
@@ -374,6 +377,7 @@
             [groups]="group._cachedSelection.groups"
             (graphResize)="storeSubDagDims($event, group)"
             [resolveReference]="resolveReference"
+            [visible]="isGroupExpanded(group)"
           ></ai-dag-raw>
         </ng-container>
       </footer>

src/app/directed_acyclic_graph_raw.ts

@@ -131,6 +131,20 @@ function setGroupSizeProps(
   return Object.assign(expandedGroup, {width, height, padY});
 }
 
+/**
+ * Sets the selection state of a group.
+ * _cachedSelection is a property that should be refactored to be clearer. It is
+ * used to determine which node/group is rendered within an iteration group.
+ */
+export function setEnhancedGroupSelection(
+    group: DagGroup,
+    selection: DagGroup|DagNode,
+) {
+  const enhancedGroup = group as EnhancedDagGroup;
+  enhancedGroup._cachedSelection = selection;
+  enhancedGroup.selectedLoopId = selection?.id;
+}
+
 /** Dimension config for a raw DAG */
 export interface GraphDims {
   width: number;
@@ -284,6 +298,12 @@ export class DagRaw implements DoCheck, OnInit, OnDestroy {
 
   @Input() loading = false;
 
+  /**
+   * If set to false, the contents of the DAG will not be rendered to optimize
+   * for performance.
+   */
+  @Input() visible = true;
+
   @Input('nodes')
   set nodes(nodes: DagNode[]) {
     // Avoid pointer/reference stability, so that angular will pick up the
@@ -471,7 +491,7 @@ export class DagRaw implements DoCheck, OnInit, OnDestroy {
       group: DagGroup, iterationNode: GroupIterationRecord['iterationNode']) {
     (group as any)._cachedSelection = iterationNode;
     if (group.selectedLoopId !== iterationNode.id) {
-      group.selectedLoopId = iterationNode.id;
+      setEnhancedGroupSelection(group, iterationNode);
       const iter = {path: this.dagPath, group, iterationNode};
       this.stateService?.setIterationChange(iter);
     }
@@ -938,6 +958,7 @@ export class DagRaw implements DoCheck, OnInit, OnDestroy {
       this.stateService?.setGroupExpandToggled({groupId: id, isExpanded: true});
     }
     this.updateGraphLayout();
+    this.detectChanges();
     return beforeCt !== this.expandedGroups.size;
   }
 

src/app/test_resources/demo_page_harness.ts

@@ -17,6 +17,7 @@
 
 import {ComponentHarness} from '@angular/cdk/testing';
 import {MatNativeSelectHarness} from '@angular/material/input/testing';
+import {MatSelectHarness} from '@angular/material/select/testing';
 
 /** Test harness for the Demo Page. */
 export class DemoPageHarness extends ComponentHarness {
@@ -31,4 +32,8 @@ export class DemoPageHarness extends ComponentHarness {
     return this.locatorFor(
         MatNativeSelectHarness.with({selector: '#color-theme-select'}))();
   }
+
+  getNodeSelectInput(): Promise<MatSelectHarness> {
+    return this.locatorFor(MatSelectHarness.with({selector: '#node-select'}))();
+  }
 }