diff --git a/packages/mermaid/src/rendering-util/layout-algorithms/ipsepCola/assignInitialPositioning.spec.ts b/packages/mermaid/src/rendering-util/layout-algorithms/ipsepCola/assignInitialPositioning.spec.ts
index b61b68f43..3076d57fe 100644
--- a/packages/mermaid/src/rendering-util/layout-algorithms/ipsepCola/assignInitialPositioning.spec.ts
+++ b/packages/mermaid/src/rendering-util/layout-algorithms/ipsepCola/assignInitialPositioning.spec.ts
@@ -1,12 +1,9 @@
 import { FlowDB } from '../../../diagrams/flowchart/flowDb.js';
 import flow from '../../../diagrams/flowchart/parser/flowParser.js';
-import type { D3Selection } from '../../../types.ts';
-import { createGraphWithElements } from '../../createGraph.js';
 import type { Node } from '../../types.ts';
 import { assignInitialPositions } from './assignInitialPositions.js';
 import { layerAssignment } from './layerAssignment.js';
 import { assignNodeOrder } from './nodeOrdering.js';
-import * as d3 from 'd3';
 
 describe('assignInitialPositioning', () => {
   beforeEach(function () {
@@ -24,10 +21,6 @@ describe('assignInitialPositioning', () => {
     await flow.parse(flowchart);
     const layoutData = flow.parser.yy.getData();
 
-    const svg: d3.Selection<SVGSVGElement, unknown, HTMLElement, any> = d3.select('svg');
-    const element = svg.select('g') as unknown as D3Selection<SVGElement>;
-    const graph = createGraphWithElements(element, layoutData);
-
     // Call Layer Assignment for the graph
     layerAssignment(layoutData);
 
@@ -67,10 +60,6 @@ describe('assignInitialPositioning', () => {
     await flow.parse(flowchart);
     const layoutData = flow.parser.yy.getData();
 
-    const svg: d3.Selection<SVGSVGElement, unknown, HTMLElement, any> = d3.select('svg');
-    const element = svg.select('g') as unknown as D3Selection<SVGElement>;
-    const graph = createGraphWithElements(element, layoutData);
-
     // Call Layer Assignment for the graph
     layerAssignment(layoutData);
 
@@ -110,10 +99,6 @@ describe('assignInitialPositioning', () => {
     await flow.parse(flowchart);
     const layoutData = flow.parser.yy.getData();
 
-    const svg: d3.Selection<SVGSVGElement, unknown, HTMLElement, any> = d3.select('svg');
-    const element = svg.select('g') as unknown as D3Selection<SVGElement>;
-    const graph = createGraphWithElements(element, layoutData);
-
     // Call Layer Assignment for the graph
     layerAssignment(layoutData);
 
@@ -160,10 +145,6 @@ describe('assignInitialPositioning', () => {
     await flow.parse(flowchart);
     const layoutData = flow.parser.yy.getData();
 
-    const svg: d3.Selection<SVGSVGElement, unknown, HTMLElement, any> = d3.select('svg');
-    const element = svg.select('g') as unknown as D3Selection<SVGElement>;
-    const graph = createGraphWithElements(element, layoutData);
-
     // Call Layer Assignment for the graph
     layerAssignment(layoutData);
 
@@ -219,10 +200,6 @@ describe('assignInitialPositioning', () => {
     await flow.parse(flowchart);
     const layoutData = flow.parser.yy.getData();
 
-    const svg: d3.Selection<SVGSVGElement, unknown, HTMLElement, any> = d3.select('svg');
-    const element = svg.select('g') as unknown as D3Selection<SVGElement>;
-    const graph = createGraphWithElements(element, layoutData);
-
     // Call Layer Assignment for the graph
     layerAssignment(layoutData);
 
diff --git a/packages/mermaid/src/rendering-util/layout-algorithms/ipsepCola/layerAssignment.spec.ts b/packages/mermaid/src/rendering-util/layout-algorithms/ipsepCola/layerAssignment.spec.ts
index 09f7dc0f8..47aa40901 100644
--- a/packages/mermaid/src/rendering-util/layout-algorithms/ipsepCola/layerAssignment.spec.ts
+++ b/packages/mermaid/src/rendering-util/layout-algorithms/ipsepCola/layerAssignment.spec.ts
@@ -1,9 +1,6 @@
-import * as d3 from 'd3';
 import { FlowDB } from '../../../diagrams/flowchart/flowDb.js';
 import flow from '../../../diagrams/flowchart/parser/flowParser.js';
-import type { D3Selection } from '../../../types.ts';
 import type { Node } from '../../types.ts';
-import { createGraphWithElements } from '../../createGraph.js';
 import { layerAssignment } from './layerAssignment.js';
 
 describe('layerAssignment', () => {
@@ -21,9 +18,6 @@ describe('layerAssignment', () => {
     // Get layout data from flowDb
     await flow.parse(flowchart);
     const layoutData = flow.parser.yy.getData();
-    const svg: d3.Selection<SVGSVGElement, unknown, HTMLElement, any> = d3.select('svg');
-    const element = svg.select('g') as unknown as D3Selection<SVGElement>;
-    const graph = createGraphWithElements(element, layoutData);
 
     // Call Layer Assignment for the graph
     layerAssignment(layoutData);
@@ -49,10 +43,6 @@ describe('layerAssignment', () => {
     await flow.parse(flowchart);
     const layoutData = flow.parser.yy.getData();
 
-    const svg: d3.Selection<SVGSVGElement, unknown, HTMLElement, any> = d3.select('svg');
-    const element = svg.select('g') as unknown as D3Selection<SVGElement>;
-    const graph = createGraphWithElements(element, layoutData);
-
     // Call Layer Assignment for the graph
     layerAssignment(layoutData);
 
@@ -79,10 +69,6 @@ describe('layerAssignment', () => {
     await flow.parse(flowchart);
     const layoutData = flow.parser.yy.getData();
 
-    const svg: d3.Selection<SVGSVGElement, unknown, HTMLElement, any> = d3.select('svg');
-    const element = svg.select('g') as unknown as D3Selection<SVGElement>;
-    const graph = createGraphWithElements(element, layoutData);
-
     // Call Layer Assignment for the graph
     layerAssignment(layoutData);
 
@@ -111,10 +97,6 @@ describe('layerAssignment', () => {
     await flow.parse(flowchart);
     const layoutData = flow.parser.yy.getData();
 
-    const svg: d3.Selection<SVGSVGElement, unknown, HTMLElement, any> = d3.select('svg');
-    const element = svg.select('g') as unknown as D3Selection<SVGElement>;
-    const graph = createGraphWithElements(element, layoutData);
-
     // Call Layer Assignment for the graph
     layerAssignment(layoutData);
 
@@ -147,10 +129,6 @@ describe('layerAssignment', () => {
     await flow.parse(flowchart);
     const layoutData = flow.parser.yy.getData();
 
-    const svg: d3.Selection<SVGSVGElement, unknown, HTMLElement, any> = d3.select('svg');
-    const element = svg.select('g') as unknown as D3Selection<SVGElement>;
-    const graph = createGraphWithElements(element, layoutData);
-
     // Call Layer Assignment for the graph
     layerAssignment(layoutData);
 
@@ -163,4 +141,21 @@ describe('layerAssignment', () => {
     expect(layoutData.nodes.find((node: Node) => node.id === 'd').layer).toEqual(2);
     expect(layoutData.nodes.find((node: Node) => node.id === 'P').layer).toEqual(3);
   });
+  it('should correctly assign the layers to node', async () => {
+    const flowchart = `
+          flowchart LR
+            A --|Test Label|--> B --> C
+        `;
+
+    // Get layout data from flowDb
+    await flow.parse(flowchart);
+    const layoutData = flow.parser.yy.getData();
+
+    // Call Layer Assignment for the graph
+    layerAssignment(layoutData);
+
+    expect(layoutData.nodes.find((node: Node) => node.id === 'A').layer).toEqual(1);
+    expect(layoutData.nodes.find((node: Node) => node.id === 'B').layer).toEqual(2);
+    expect(layoutData.nodes.find((node: Node) => node.id === 'C').layer).toEqual(3);
+  });
 });
diff --git a/packages/mermaid/src/rendering-util/layout-algorithms/ipsepCola/nodeOrdering.spec.ts b/packages/mermaid/src/rendering-util/layout-algorithms/ipsepCola/nodeOrdering.spec.ts
index d0f83bc90..e6a30295a 100644
--- a/packages/mermaid/src/rendering-util/layout-algorithms/ipsepCola/nodeOrdering.spec.ts
+++ b/packages/mermaid/src/rendering-util/layout-algorithms/ipsepCola/nodeOrdering.spec.ts
@@ -1,9 +1,6 @@
 import { FlowDB } from '../../../diagrams/flowchart/flowDb.js';
 import flow from '../../../diagrams/flowchart/parser/flowParser.js';
-import type { D3Selection } from '../../../types.ts';
-import { createGraphWithElements } from '../../createGraph.js';
 import type { Node } from '../../types.ts';
-import * as d3 from 'd3';
 import { layerAssignment } from './layerAssignment.js';
 import { assignNodeOrder } from './nodeOrdering.js';
 
@@ -23,10 +20,6 @@ describe('nodeOrdering', () => {
     await flow.parse(flowchart);
     const layoutData = flow.parser.yy.getData();
 
-    const svg: d3.Selection<SVGSVGElement, unknown, HTMLElement, any> = d3.select('svg');
-    const element = svg.select('g') as unknown as D3Selection<SVGElement>;
-    const graph = createGraphWithElements(element, layoutData);
-
     // Call Layer Assignment for the graph
     layerAssignment(layoutData);
 
@@ -54,10 +47,6 @@ describe('nodeOrdering', () => {
     await flow.parse(flowchart);
     const layoutData = flow.parser.yy.getData();
 
-    const svg: d3.Selection<SVGSVGElement, unknown, HTMLElement, any> = d3.select('svg');
-    const element = svg.select('g') as unknown as D3Selection<SVGElement>;
-    const graph = createGraphWithElements(element, layoutData);
-
     // Call Layer Assignment for the graph
     layerAssignment(layoutData);
 
@@ -87,10 +76,6 @@ describe('nodeOrdering', () => {
     await flow.parse(flowchart);
     const layoutData = flow.parser.yy.getData();
 
-    const svg: d3.Selection<SVGSVGElement, unknown, HTMLElement, any> = d3.select('svg');
-    const element = svg.select('g') as unknown as D3Selection<SVGElement>;
-    const graph = createGraphWithElements(element, layoutData);
-
     // Call Layer Assignment for the graph
     layerAssignment(layoutData);
 
@@ -121,9 +106,6 @@ describe('nodeOrdering', () => {
     await flow.parse(flowchart);
     const layoutData = flow.parser.yy.getData();
 
-    const svg: d3.Selection<SVGSVGElement, unknown, HTMLElement, any> = d3.select('svg');
-    const element = svg.select('g') as unknown as D3Selection<SVGElement>;
-    const graph = createGraphWithElements(element, layoutData);
     // Call Layer Assignment for the graph
     layerAssignment(layoutData);
 
@@ -159,10 +141,6 @@ describe('nodeOrdering', () => {
     await flow.parse(flowchart);
     const layoutData = flow.parser.yy.getData();
 
-    const svg: d3.Selection<SVGSVGElement, unknown, HTMLElement, any> = d3.select('svg');
-    const element = svg.select('g') as unknown as D3Selection<SVGElement>;
-    const graph = createGraphWithElements(element, layoutData);
-
     // Call Layer Assignment for the graph
     layerAssignment(layoutData);