diff --git a/libs/cde-visualization/src/lib/cde-visualization/cde-visualization.component.spec.ts b/libs/cde-visualization/src/lib/cde-visualization/cde-visualization.component.spec.ts
index 98cefb6d7..3c488f069 100644
--- a/libs/cde-visualization/src/lib/cde-visualization/cde-visualization.component.spec.ts
+++ b/libs/cde-visualization/src/lib/cde-visualization/cde-visualization.component.spec.ts
@@ -4,8 +4,10 @@ import userEvent from '@testing-library/user-event';
 import { mockDeep } from 'jest-mock-extended';
 import embed, { Result } from 'vega-embed';
 
-import { provideScrolling } from '@hra-ui/design-system/scrolling';
+import { TestbedHarnessEnvironment } from '@angular/cdk/testing/testbed';
+import { MatMenuHarness } from '@angular/material/menu/testing';
 import { rgbToHex } from '@hra-ui/design-system/color-picker';
+import { provideScrolling } from '@hra-ui/design-system/scrolling';
 import { ColorMapEntry, DEFAULT_COLOR_MAP_KEY, DEFAULT_COLOR_MAP_VALUE_KEY } from '../models/color-map';
 import { EdgeEntry } from '../models/edge';
 import { DEFAULT_NODE_TARGET_KEY, DEFAULT_NODE_TARGET_VALUE, NodeEntry } from '../models/node';
@@ -144,9 +146,7 @@ describe('CdeVisualizationComponent', () => {
   });
 
   it('should update color map when downloadColorMap is called', async () => {
-    const {
-      fixture: { componentInstance: instance },
-    } = await setup({
+    const { fixture } = await setup({
       componentInputs: {
         ...sampleData,
         nodes: sampleNodes,
@@ -154,6 +154,7 @@ describe('CdeVisualizationComponent', () => {
       },
     });
 
+    const instance = fixture.componentInstance;
     const processedColorMap = instance
       .cellTypesAsColorMap()
       .map((entry) => ({ ...entry, [instance.colorMapValueKey()]: rgbToHex(entry[instance.colorMapValueKey()]) }));
@@ -161,8 +162,10 @@ describe('CdeVisualizationComponent', () => {
     const fileSaver = TestBed.inject(FileSaverService);
     const fileSaveSpy = jest.spyOn(fileSaver, 'saveCsv').mockReturnValue(undefined);
 
-    const downloadColorMapButton = screen.getByText('CSV');
-    await userEvent.click(downloadColorMapButton);
+    const loader = TestbedHarnessEnvironment.loader(fixture);
+    const menu = await loader.getHarness(MatMenuHarness);
+
+    await menu.clickItem({ text: /Download/ }, { text: /Cell Color Map CSV/ });
     expect(fileSaveSpy).toHaveBeenCalledWith(processedColorMap, 'color-map.csv');
   });
 
diff --git a/libs/cde-visualization/src/lib/components/cell-types/cell-types.component.html b/libs/cde-visualization/src/lib/components/cell-types/cell-types.component.html
index 64fd3f807..ffcd0c7e9 100644
--- a/libs/cde-visualization/src/lib/components/cell-types/cell-types.component.html
+++ b/libs/cde-visualization/src/lib/components/cell-types/cell-types.component.html
@@ -1,6 +1,6 @@
 <span class="cell-types-label">
   Cell Types
-  <button mat-icon-button [matMenuTriggerFor]="menu" hraIconButtonSize="medium">
+  <button mat-icon-button [matMenuTriggerFor]="menu" hraIconButtonSize="medium" data-testid="trigger">
     <mat-icon>more_vert</mat-icon>
   </button>
   <mat-menu #menu="matMenu">
diff --git a/libs/cde-visualization/src/lib/components/node-dist-visualization/node-dist-visualization.component.spec.ts b/libs/cde-visualization/src/lib/components/node-dist-visualization/node-dist-visualization.component.spec.ts
index 7cef99e51..90be09dbf 100644
--- a/libs/cde-visualization/src/lib/components/node-dist-visualization/node-dist-visualization.component.spec.ts
+++ b/libs/cde-visualization/src/lib/components/node-dist-visualization/node-dist-visualization.component.spec.ts
@@ -54,9 +54,9 @@ describe('NodeDistVisualizationComponent', () => {
     [2, 0, 4, 3, 4, 5, 6],
   ];
   const sampleCellTypes: CellTypeEntry[] = [
-    { name: 'a', count: 2, color: [0, 0, 0] },
-    { name: 'b', count: 4, color: [0, 1, 2] },
-    { name: 'c', count: 6, color: [0, 1, 3] },
+    { name: 'a', count: 2, color: [0, 0, 0], outgoingEdgeCount: 0 },
+    { name: 'b', count: 4, color: [0, 1, 2], outgoingEdgeCount: 0 },
+    { name: 'c', count: 6, color: [0, 1, 3], outgoingEdgeCount: 0 },
   ];
   const sampleCellTypesSelection: string[] = [sampleCellTypes[0].name, sampleCellTypes[1].name];