prowler/ui/lib/attack-paths/layout.test.ts

import { Position } from "@xyflow/react";
import { describe, expect, it } from "vitest";

import { layoutWithDagre } from "@/lib/attack-paths/layout";
import type { GraphEdge, GraphNode } from "@/types/attack-paths";

const findingNode: GraphNode = {
  id: "finding-1",
  labels: ["ProwlerFinding"],
  properties: { check_title: "Open S3 bucket", severity: "high" },
};

const resourceNode: GraphNode = {
  id: "resource-1",
  labels: ["S3Bucket"],
  properties: { name: "bucket-1" },
};

const guardDutyNode: GraphNode = {
  id: "guard-duty-1",
  labels: ["GuardDutyFinding"],
  properties: { title: "Port probe", severity: "high" },
};

const inspectorNode: GraphNode = {
  id: "inspector-1",
  labels: ["AWSInspectorFinding"],
  properties: { title: "Package vulnerability", severity: "high" },
};

const internetNode: GraphNode = {
  id: "internet-1",
  labels: ["Internet"],
  properties: {},
};

describe("layoutWithDagre", () => {
  it("returns empty arrays for empty input", () => {
    const result = layoutWithDagre([], []);
    expect(result.rfNodes).toEqual([]);
    expect(result.rfEdges).toEqual([]);
  });
  it("assigns node types and dimensions from labels", () => {
    const { rfNodes } = layoutWithDagre(
      [findingNode, resourceNode, internetNode],
      [],
    );

    const byId = new Map(rfNodes.map((n) => [n.id, n]));

    expect(byId.get("finding-1")).toMatchObject({
      type: "finding",
      width: 150,
      height: 124,
    });
    expect(byId.get("resource-1")).toMatchObject({
      type: "resource",
      width: 136,
      height: 124,
    });
    expect(byId.get("internet-1")).toMatchObject({
      type: "internet",
      width: 80,
      height: 80,
    });
  });

  it("treats cloud-provider finding resources as resource nodes", () => {
    const { rfNodes } = layoutWithDagre(
      [findingNode, guardDutyNode, inspectorNode],
      [],
    );

    const byId = new Map(rfNodes.map((n) => [n.id, n]));

    expect(byId.get("finding-1")?.type).toBe("finding");
    expect(byId.get("guard-duty-1")).toMatchObject({
      type: "resource",
      width: 136,
      height: 124,
    });
    expect(byId.get("inspector-1")?.type).toBe("resource");
  });

  it("does not animate edges that only touch cloud-provider finding resources", () => {
    const { rfEdges } = layoutWithDagre(
      [resourceNode, guardDutyNode],
      [
        {
          id: "e1",
          source: "guard-duty-1",
          target: "resource-1",
          type: "AFFECTS",
        },
      ],
    );

    expect(rfEdges[0]).toMatchObject({
      animated: false,
      className: "resource-edge",
    });
  });

  it("is deterministic: same input produces equal output across runs", () => {
    const nodes = [findingNode, resourceNode];
    const edges: GraphEdge[] = [
      {
        id: "e1",
        source: "resource-1",
        target: "finding-1",
        type: "HAS_FINDING",
      },
    ];

    const a = layoutWithDagre(nodes, edges);
    const b = layoutWithDagre(nodes, edges);

    expect(a).toEqual(b);
  });

  it("places connected children to the right and stacks siblings within the horizontal rank", () => {
    const rootNode: GraphNode = {
      id: "root",
      labels: ["AWSAccount"],
      properties: { name: "account" },
    };
    const siblingNodes: GraphNode[] = [
      {
        id: "bucket",
        labels: ["S3Bucket"],
        properties: { name: "bucket" },
      },
      {
        id: "lambda",
        labels: ["AWSLambda"],
        properties: { name: "function" },
      },
      {
        id: "database",
        labels: ["RDSInstance"],
        properties: { name: "database" },
      },
    ];

    const { rfNodes } = layoutWithDagre(
      [rootNode, ...siblingNodes],
      siblingNodes.map((node) => ({
        id: `root-${node.id}`,
        source: "root",
        target: node.id,
        type: "CONNECTS_TO",
      })),
    );

    const rootPosition = rfNodes.find(
      (candidate) => candidate.id === "root",
    )?.position;
    const siblingPositions = siblingNodes.map((node) => {
      const rfNode = rfNodes.find((candidate) => candidate.id === node.id);

      expect(rfNode).toBeDefined();

      return rfNode?.position ?? { x: 0, y: 0 };
    });

    const xSpread =
      Math.max(...siblingPositions.map((position) => position.x)) -
      Math.min(...siblingPositions.map((position) => position.x));
    const ySpread =
      Math.max(...siblingPositions.map((position) => position.y)) -
      Math.min(...siblingPositions.map((position) => position.y));

    expect(rootPosition).toBeDefined();
    siblingPositions.forEach((position) => {
      expect(position.x).toBeGreaterThan(rootPosition?.x ?? 0);
    });
    expect(ySpread).toBeGreaterThan(xSpread);
  });

  it("connects edges through right and left node sides for horizontal layout", () => {
    const { rfNodes } = layoutWithDagre(
      [findingNode, resourceNode],
      [
        {
          id: "e1",
          source: "resource-1",
          target: "finding-1",
          type: "HAS_FINDING",
        },
      ],
    );

    rfNodes.forEach((node) => {
      expect(node.sourcePosition).toBe(Position.Right);
      expect(node.targetPosition).toBe(Position.Left);
    });
  });

  it("offsets dagre center positions by half of the node dimensions (top-left)", () => {
    const { rfNodes } = layoutWithDagre([findingNode, resourceNode], []);

    rfNodes.forEach((node) => {
      expect(Number.isFinite(node.position.x)).toBe(true);
      expect(Number.isFinite(node.position.y)).toBe(true);
    });

    // Different node types must end up with different sizes — confirms the
    // dimension-aware offset is wired up.
    const findingDims = rfNodes.find((n) => n.id === "finding-1");
    const resourceDims = rfNodes.find((n) => n.id === "resource-1");
    expect(findingDims?.width).not.toEqual(resourceDims?.width);
  });

  it("reverses container relationships while preserving original endpoints in edge data", () => {
    const containerNode: GraphNode = {
      id: "container",
      labels: ["AWSAccount"],
      properties: { name: "acct" },
    };
    const childNode: GraphNode = {
      id: "child",
      labels: ["S3Bucket"],
      properties: { name: "bucket" },
    };

    const { rfEdges } = layoutWithDagre(
      [containerNode, childNode],
      [
        {
          id: "e1",
          source: "container",
          target: "child",
          type: "RUNS_IN",
        },
      ],
    );

    expect(rfEdges).toHaveLength(1);
    expect(rfEdges[0]).toMatchObject({
      source: "child",
      target: "container",
      data: { originalSource: "container", originalTarget: "child" },
    });
  });

  it("animates edges that touch a finding node and tags them with finding-edge", () => {
    const { rfEdges } = layoutWithDagre(
      [findingNode, resourceNode, internetNode],
      [
        {
          id: "e1",
          source: "resource-1",
          target: "finding-1",
          type: "HAS_FINDING",
        },
        {
          id: "e2",
          source: "internet-1",
          target: "resource-1",
          type: "CONNECTS_TO",
        },
      ],
    );

    const findingEdge = rfEdges.find(
      (e) => e.source === "resource-1" && e.target === "finding-1",
    );
    const plainEdge = rfEdges.find(
      (e) => e.source === "internet-1" && e.target === "resource-1",
    );

    expect(findingEdge).toMatchObject({
      animated: true,
      className: "finding-edge",
    });
    expect(plainEdge).toMatchObject({
      animated: false,
      className: "resource-edge",
    });
  });

  it("preserves the original edge id when the graph has a single edge", () => {
    const { rfEdges } = layoutWithDagre(
      [findingNode, resourceNode],
      [
        {
          id: "ignored-by-rf",
          source: "resource-1",
          target: "finding-1",
          type: "HAS_FINDING",
        },
      ],
    );

    expect(rfEdges[0]?.id).toBe("ignored-by-rf");
  });

  it("preserves parallel edges between the same nodes with unique ids", () => {
    const { rfEdges } = layoutWithDagre(
      [resourceNode, findingNode],
      [
        {
          id: "edge-1",
          source: "resource-1",
          target: "finding-1",
          type: "HAS_FINDING",
        },
        {
          id: "edge-2",
          source: "resource-1",
          target: "finding-1",
          type: "HAS_FINDING",
        },
      ],
    );

    expect(rfEdges).toHaveLength(2);
    expect(rfEdges.map((edge) => edge.id)).toEqual(["edge-1", "edge-2"]);
    expect(rfEdges.every((edge) => edge.source === "resource-1")).toBe(true);
    expect(rfEdges.every((edge) => edge.target === "finding-1")).toBe(true);
  });
});