Fix multi label/class support in neighbor counts

packages/graph-explorer/src/connector/gremlin/neighborCounts.test.ts

@@ -1,4 +1,4 @@
-import { createVertexId } from "@/core";
+import { createVertexId, EntityRawId } from "@/core";
 import { neighborCounts } from "./neighborCounts";
 import { query } from "@/utils";
 import { NeighborCount } from "../useGEFetchTypes";
@@ -7,6 +7,7 @@ import {
   createGremlinResponse,
   createRandomVertexId,
 } from "@/utils/testing";
+import { GMap } from "./types";
 
 describe("neighborCounts", () => {
   it("should return empty for an empty request", async () => {
@@ -42,6 +43,7 @@ describe("neighborCounts", () => {
       counts: {
         label1: 3,
         label2: 9,
+        label3: 9,
       },
     };
     const response = createResponse({
@@ -116,18 +118,86 @@ describe("neighborCounts", () => {
        )
     `);
   });
+
+  it("should handle error response", async () => {
+    const mockFetch = vi.fn().mockResolvedValue({
+      code: 500,
+      detailedMessage: "Internal server error occurred",
+    });
+
+    await expect(
+      neighborCounts(mockFetch, {
+        vertexIds: [createVertexId("123")],
+      })
+    ).rejects.toThrow("Internal server error occurred");
+  });
+
+  it("should handle empty response data", async () => {
+    const mockFetch = vi
+      .fn()
+      .mockResolvedValue(createGremlinResponse(createGMap({})));
+
+    const result = await neighborCounts(mockFetch, {
+      vertexIds: [createVertexId("123")],
+    });
+
+    expect(result.counts).toEqual([]);
+  });
+
+  it("should handle vertex with zero neighbors", async () => {
+    const vertexId = createRandomVertexId();
+    const expected: NeighborCount = {
+      vertexId,
+      totalCount: 0,
+      counts: {},
+    };
+    const response = createResponse(expected);
+    const mockFetch = vi.fn().mockResolvedValue(response);
+
+    const result = await neighborCounts(mockFetch, {
+      vertexIds: [vertexId],
+    });
+
+    expect(result.counts).toEqual([expected]);
+  });
+
+  it("should handle mixed vertex ID types", async () => {
+    const stringId = createVertexId("string-id");
+    const numberId = createVertexId(42);
+    const expected1: NeighborCount = {
+      vertexId: stringId,
+      totalCount: 5,
+      counts: { type1: 5 },
+    };
+    const expected2: NeighborCount = {
+      vertexId: numberId, // Gremlin converts number IDs to strings
+      totalCount: 3,
+      counts: { type2: 3 },
+    };
+    const response = createResponse(expected1, expected2);
+    const mockFetch = vi.fn().mockResolvedValue(response);
+
+    const result = await neighborCounts(mockFetch, {
+      vertexIds: [stringId, numberId],
+    });
+
+    expect(result.counts).toHaveLength(2);
+    expect(result.counts).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining(expected1),
+        expect.objectContaining(expected2),
+      ])
+    );
+  });
 });
 
 function createResponse(...counts: NeighborCount[]) {
   return createGremlinResponse(
     createGMap(
-      counts.reduce(
-        (prev, curr) => {
-          prev[String(curr.vertexId)] = createGMap(curr.counts);
-          return prev;
-        },
-        {} as Record<string, any>
-      )
+      counts.reduce((prev, curr) => {
+        prev.set(curr.vertexId, createGMap(curr.counts));
+        return prev;
+      }, new Map<EntityRawId, GMap>())
     )
   );
 }

packages/graph-explorer/src/connector/gremlin/neighborCounts.ts

@@ -83,21 +83,20 @@ export async function neighborCounts(
         }))
         .reduce(
           (acc, curr) => {
-            // TODO: In a future set of changes we should pass the full lsit of types
-            // up to the UI so that it can list them out properly, but since this is a
-            // rather large change I am defering that work.
-            const type = curr.type.split("::")[0] ?? "";
-            acc[type] = curr.count;
+            const types = curr.type.split("::");
+            for (const type of types) {
+              acc[type] = (acc[type] ?? 0) + curr.count;
+            }
             return acc;
           },
           {} as Record<string, number>
         );
 
-      // Sum up all the type counts
-      const totalCount = Object.values(countsByType).reduce(
-        (acc, curr) => acc + curr,
-        0
-      );
+      // Total up the unique neighbors
+      const totalCount = countsByTypeMap
+        .values()
+        .map(gValue => gValue["@value"])
+        .reduce((acc, curr) => acc + curr, 0);
 
       return {
         vertexId: createVertexId(extractRawId(key)),

packages/graph-explorer/src/connector/openCypher/neighborCounts.test.ts

@@ -49,13 +49,14 @@ describe("neighborCounts", () => {
     expect(result.counts).toEqual([expected]);
   });
 
-  it("should return neighbor counts", async () => {
+  it("should return neighbor counts with multiple labels", async () => {
     const expected: NeighborCount = {
       vertexId: createRandomVertexId(),
       totalCount: 12,
       counts: {
         label1: 3,
         label2: 9,
+        label3: 9,
       },
     };
     const response = createResponse({
@@ -96,6 +97,134 @@ describe("neighborCounts", () => {
     });
     expect(result.counts).toEqual([expected1, expected2]);
   });
+
+  it("should handle error response", async () => {
+    const mockFetch = vi.fn().mockResolvedValue({
+      code: 500,
+      detailedMessage: "openCypher query failed",
+    });
+
+    await expect(
+      neighborCounts(mockFetch, {
+        vertexIds: [createVertexId("123")],
+      })
+    ).rejects.toThrow("openCypher query failed");
+  });
+
+  it("should handle empty results", async () => {
+    const mockFetch = vi.fn().mockResolvedValue({ results: [] });
+
+    const result = await neighborCounts(mockFetch, {
+      vertexIds: [createVertexId("123")],
+    });
+
+    expect(result.counts).toEqual([]);
+  });
+
+  it("should handle vertex with zero neighbors", async () => {
+    const vertexId = createRandomVertexId();
+    const response = {
+      results: [
+        {
+          id: String(vertexId),
+          counts: [],
+        },
+      ],
+    };
+    const mockFetch = vi.fn().mockResolvedValue(response);
+
+    const result = await neighborCounts(mockFetch, {
+      vertexIds: [vertexId],
+    });
+
+    expect(result.counts).toEqual([
+      {
+        vertexId,
+        totalCount: 0,
+        counts: {},
+      },
+    ]);
+  });
+
+  it("should handle malformed result data", async () => {
+    const vertexId = createRandomVertexId();
+    const response = {
+      results: [
+        {
+          id: String(vertexId),
+          // Missing counts property
+        },
+      ],
+    };
+    const mockFetch = vi.fn().mockResolvedValue(response);
+
+    const result = await neighborCounts(mockFetch, {
+      vertexIds: [vertexId],
+    });
+
+    expect(result.counts).toEqual([]);
+  });
+
+  it("should handle single label neighbors", async () => {
+    const vertexId = createRandomVertexId();
+    const expected: NeighborCount = {
+      vertexId,
+      totalCount: 5,
+      counts: {
+        Person: 5,
+      },
+    };
+    const response = {
+      results: [
+        {
+          id: String(vertexId),
+          counts: [
+            {
+              label: ["Person"],
+              count: 5,
+            },
+          ],
+        },
+      ],
+    };
+    const mockFetch = vi.fn().mockResolvedValue(response);
+
+    const result = await neighborCounts(mockFetch, {
+      vertexIds: [vertexId],
+    });
+
+    expect(result.counts).toEqual([expected]);
+  });
+
+  it("should handle empty label arrays", async () => {
+    const vertexId = createRandomVertexId();
+    const response = {
+      results: [
+        {
+          id: String(vertexId),
+          counts: [
+            {
+              label: [],
+              count: 3,
+            },
+          ],
+        },
+      ],
+    };
+    const mockFetch = vi.fn().mockResolvedValue(response);
+
+    const result = await neighborCounts(mockFetch, {
+      vertexIds: [vertexId],
+    });
+
+    expect(result.counts).toEqual([
+      {
+        vertexId,
+        totalCount: 3,
+        counts: {},
+      },
+    ]);
+  });
 });
 
 function createResponse(...counts: NeighborCount[]) {

packages/graph-explorer/src/connector/openCypher/neighborCounts.ts

@@ -53,14 +53,16 @@ export async function neighborCounts(
     }
 
     const vertexId = createVertexId(rawId);
+
+    // Total up neighbors by type (might be more than unique neighbors)
     const countsByType: Record<string, number> = {};
     for (const count of rawCounts) {
-      // TODO: In a future set of changes we should pass the full lsit of types
-      // up to the UI so that it can list them out properly, but since this is a
-      // rather large change I am defering that work.
-      const type = count.label[0] ?? "";
-      countsByType[type] = count.count;
+      for (const type of count.label) {
+        countsByType[type] = (countsByType[type] ?? 0) + count.count;
+      }
     }
+
+    // Total up the unique neighbors
     const totalCount = rawCounts.reduce((acc, count) => acc + count.count, 0);
 
     counts.push({