Fix gremlin counts

Author: kmcginnes

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)),