rivetkit 2.0.2 → 2.0.3

package/src/driver-test-suite/tests/actor-conn-state.ts

@@ -0,0 +1,249 @@
+import { describe, expect, test, vi } from "vitest";
+import type { DriverTestConfig } from "../mod";
+import { setupDriverTest } from "../utils";
+
+export function runActorConnStateTests(driverTestConfig: DriverTestConfig) {
+	describe("Actor Connection State Tests", () => {
+		describe("Connection State Initialization", () => {
+			test("should retrieve connection state", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Connect to the actor
+				const connection = client.connStateActor.getOrCreate().connect();
+
+				// Get the connection state
+				const connState = await connection.getConnectionState();
+
+				// Verify the connection state structure
+				expect(connState.id).toBeDefined();
+				expect(connState.username).toBeDefined();
+				expect(connState.role).toBeDefined();
+				expect(connState.counter).toBeDefined();
+				expect(connState.createdAt).toBeDefined();
+
+				// Clean up
+				await connection.dispose();
+			});
+
+			test("should initialize connection state with custom parameters", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Connect with custom parameters
+				const connection = client.connStateActor
+					.getOrCreate([], {
+						params: {
+							username: "testuser",
+							role: "admin",
+						},
+					})
+					.connect();
+
+				// Get the connection state
+				const connState = await connection.getConnectionState();
+
+				// Verify the connection state was initialized with custom values
+				expect(connState.username).toBe("testuser");
+				expect(connState.role).toBe("admin");
+
+				// Clean up
+				await connection.dispose();
+			});
+		});
+
+		describe("Connection State Management", () => {
+			test("should maintain unique state for each connection", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Create multiple connections
+				const conn1 = client.connStateActor
+					.getOrCreate([], {
+						params: { username: "user1" },
+					})
+					.connect();
+
+				const conn2 = client.connStateActor
+					.getOrCreate([], {
+						params: { username: "user2" },
+					})
+					.connect();
+
+				// Update connection state for each connection
+				await conn1.incrementConnCounter(5);
+				await conn2.incrementConnCounter(10);
+
+				// Get state for each connection
+				const state1 = await conn1.getConnectionState();
+				const state2 = await conn2.getConnectionState();
+
+				// Verify states are separate
+				expect(state1.counter).toBe(5);
+				expect(state2.counter).toBe(10);
+				expect(state1.username).toBe("user1");
+				expect(state2.username).toBe("user2");
+
+				// Clean up
+				await conn1.dispose();
+				await conn2.dispose();
+			});
+
+			test("should track connections in shared state", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Create two connections
+				const handle = client.connStateActor.getOrCreate();
+				const conn1 = handle.connect();
+				const conn2 = handle.connect();
+
+				// HACK: Wait for both connections to successfully connect by waiting for a round trip RPC
+				await conn1.getConnectionState();
+				await conn2.getConnectionState();
+
+				// Get state1 for reference
+				const state1 = await conn1.getConnectionState();
+
+				// Get connection IDs tracked by the actor
+				const connectionIds = await conn1.getConnectionIds();
+
+				// There should be at least 2 connections tracked
+				expect(connectionIds.length).toBeGreaterThanOrEqual(2);
+
+				// Should include the ID of the first connection
+				expect(connectionIds).toContain(state1.id);
+
+				// Clean up
+				await conn1.dispose();
+				await conn2.dispose();
+			});
+
+			test("should identify different connections in the same actor", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Create two connections to the same actor
+				const handle = client.connStateActor.getOrCreate();
+				const conn1 = handle.connect();
+				const conn2 = handle.connect();
+
+				// HACK: Wait for both connections to successfully connect by waiting for a round trip RPC
+				await conn1.getConnectionState();
+				await conn2.getConnectionState();
+
+				// Get all connection states
+				const allStates = await conn1.getAllConnectionStates();
+
+				// Should have at least 2 states
+				expect(allStates.length).toBeGreaterThanOrEqual(2);
+
+				// IDs should be unique
+				const ids = allStates.map((state) => state.id);
+				const uniqueIds = [...new Set(ids)];
+				expect(uniqueIds.length).toBe(ids.length);
+
+				// Clean up
+				await conn1.dispose();
+				await conn2.dispose();
+			});
+		});
+
+		describe("Connection Lifecycle", () => {
+			test("should track connection and disconnection events", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Create a connection
+				const handle = client.connStateActor.getOrCreate();
+				const conn = handle.connect();
+
+				// Get the connection state
+				const connState = await conn.getConnectionState();
+
+				// Verify the connection is tracked
+				const connectionIds = await conn.getConnectionIds();
+				expect(connectionIds).toContain(connState.id);
+
+				// Initial disconnection count
+				const initialDisconnections = await conn.getDisconnectionCount();
+
+				// Dispose the connection
+				await conn.dispose();
+
+				// Create a new connection to check the disconnection count
+				const newConn = handle.connect();
+
+				// Verify disconnection was tracked
+				await vi.waitFor(async () => {
+					const newDisconnections = await newConn.getDisconnectionCount();
+
+					expect(newDisconnections).toBeGreaterThan(initialDisconnections);
+				});
+
+				// Clean up
+				await newConn.dispose();
+			});
+
+			test("should update connection state", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Create a connection
+				const conn = client.connStateActor.getOrCreate().connect();
+
+				// Get the initial state
+				const initialState = await conn.getConnectionState();
+				expect(initialState.username).toBe("anonymous");
+
+				// Update the connection state
+				const updatedState = await conn.updateConnection({
+					username: "newname",
+					role: "moderator",
+				});
+
+				// Verify the state was updated
+				expect(updatedState.username).toBe("newname");
+				expect(updatedState.role).toBe("moderator");
+
+				// Get the state again to verify persistence
+				const latestState = await conn.getConnectionState();
+				expect(latestState.username).toBe("newname");
+				expect(latestState.role).toBe("moderator");
+
+				// Clean up
+				await conn.dispose();
+			});
+		});
+
+		describe("Connection Communication", () => {
+			test("should send messages to specific connections", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Create two connections
+				const handle = client.connStateActor.getOrCreate();
+				const conn1 = handle.connect();
+				const conn2 = handle.connect();
+
+				// Get connection states
+				const state1 = await conn1.getConnectionState();
+				const state2 = await conn2.getConnectionState();
+
+				// Set up event listener on second connection
+				const receivedMessages: any[] = [];
+				conn2.on("directMessage", (data) => {
+					receivedMessages.push(data);
+				});
+
+				// Send message from first connection to second
+				const success = await conn1.sendToConnection(
+					state2.id,
+					"Hello from conn1",
+				);
+				expect(success).toBe(true);
+
+				// Verify message was received
+				expect(receivedMessages.length).toBe(1);
+				expect(receivedMessages[0].from).toBe(state1.id);
+				expect(receivedMessages[0].message).toBe("Hello from conn1");
+
+				// Clean up
+				await conn1.dispose();
+				await conn2.dispose();
+			});
+		});
+	});
+}

package/src/driver-test-suite/tests/actor-conn.ts

@@ -0,0 +1,349 @@
+import { describe, expect, test } from "vitest";
+import type { DriverTestConfig } from "../mod";
+import { FAKE_TIME, setupDriverTest, waitFor } from "../utils";
+
+export function runActorConnTests(driverTestConfig: DriverTestConfig) {
+	describe("Actor Connection Tests", () => {
+		describe("Connection Methods", () => {
+			test("should connect using .get().connect()", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Create actor
+				await client.counter.create(["test-get"]);
+
+				// Get a handle and connect
+				const handle = client.counter.get(["test-get"]);
+				const connection = handle.connect();
+
+				// Verify connection by performing an action
+				const count = await connection.increment(5);
+				expect(count).toBe(5);
+
+				// Clean up
+				await connection.dispose();
+			});
+
+			test("should connect using .getForId().connect()", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Create a actor first to get its ID
+				const handle = client.counter.getOrCreate(["test-get-for-id"]);
+				await handle.increment(3);
+				const actorId = await handle.resolve();
+
+				// Get a new handle using the actor ID and connect
+				const idHandle = client.counter.getForId(actorId);
+				const connection = idHandle.connect();
+
+				// Verify connection works and state is preserved
+				const count = await connection.getCount();
+				expect(count).toBe(3);
+
+				// Clean up
+				await connection.dispose();
+			});
+
+			test("should connect using .getOrCreate().connect()", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Get or create actor and connect
+				const handle = client.counter.getOrCreate(["test-get-or-create"]);
+				const connection = handle.connect();
+
+				// Verify connection works
+				const count = await connection.increment(7);
+				expect(count).toBe(7);
+
+				// Clean up
+				await connection.dispose();
+			});
+
+			test("should connect using (await create()).connect()", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Create actor and connect
+				const handle = await client.counter.create(["test-create"]);
+				const connection = handle.connect();
+
+				// Verify connection works
+				const count = await connection.increment(9);
+				expect(count).toBe(9);
+
+				// Clean up
+				await connection.dispose();
+			});
+		});
+
+		describe("Event Communication", () => {
+			test("should mix RPC calls and WebSocket events", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Create actor
+				const handle = client.counter.getOrCreate(["test-mixed-rpc-ws"]);
+				const connection = handle.connect();
+
+				// Set up event listener
+				const receivedEvents: number[] = [];
+				const receivedEventsPromise = new Promise((resolve) => {
+					connection.on("newCount", (count: number) => {
+						receivedEvents.push(count);
+						if (
+							receivedEvents.includes(1) &&
+							receivedEvents.includes(6) &&
+							receivedEvents.includes(9)
+						)
+							resolve(undefined);
+					});
+				});
+
+				// Send one RPC call over the connection to ensure it's open
+				await connection.increment(1);
+
+				// Now use stateless RPC calls through the handle (no connection)
+				// These should still trigger events that the connection receives
+				await handle.increment(5);
+				await handle.increment(3);
+
+				// Wait for all events to be received
+				await receivedEventsPromise;
+
+				// Clean up
+				await connection.dispose();
+			});
+
+			test("should receive events via broadcast", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Create actor and connect
+				const handle = client.counter.getOrCreate(["test-broadcast"]);
+				const connection = handle.connect();
+
+				// Set up event listener
+				const receivedEvents: number[] = [];
+				connection.on("newCount", (count: number) => {
+					receivedEvents.push(count);
+				});
+
+				// Trigger broadcast events
+				await connection.increment(5);
+				await connection.increment(3);
+
+				// Verify events were received
+				expect(receivedEvents).toContain(5);
+				expect(receivedEvents).toContain(8);
+
+				// Clean up
+				await connection.dispose();
+			});
+
+			test("should handle one-time events with once()", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Create actor and connect
+				const handle = client.counter.getOrCreate(["test-once"]);
+				const connection = handle.connect();
+
+				// Set up one-time event listener
+				const receivedEvents: number[] = [];
+				connection.once("newCount", (count: number) => {
+					receivedEvents.push(count);
+				});
+
+				// Trigger multiple events, but should only receive the first one
+				await connection.increment(5);
+				await connection.increment(3);
+
+				// Verify only the first event was received
+				expect(receivedEvents).toEqual([5]);
+				expect(receivedEvents).not.toContain(8);
+
+				// Clean up
+				await connection.dispose();
+			});
+
+			test("should unsubscribe from events", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Create actor and connect
+				const handle = client.counter.getOrCreate(["test-unsubscribe"]);
+				const connection = handle.connect();
+
+				// Set up event listener with unsubscribe
+				const receivedEvents: number[] = [];
+				const unsubscribe = connection.on("newCount", (count: number) => {
+					receivedEvents.push(count);
+				});
+
+				// Trigger first event
+				await connection.increment(5);
+
+				// Unsubscribe
+				unsubscribe();
+
+				// Trigger second event, should not be received
+				await connection.increment(3);
+
+				// Verify only the first event was received
+				expect(receivedEvents).toEqual([5]);
+				expect(receivedEvents).not.toContain(8);
+
+				// Clean up
+				await connection.dispose();
+			});
+		});
+
+		describe("Connection Parameters", () => {
+			test("should pass connection parameters", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Create two connections with different params
+				const handle1 = client.counterWithParams.getOrCreate(["test-params"], {
+					params: { name: "user1" },
+				});
+				const handle2 = client.counterWithParams.getOrCreate(["test-params"], {
+					params: { name: "user2" },
+				});
+
+				const conn1 = handle1.connect();
+				const conn2 = handle2.connect();
+
+				// HACK: Call an action to wait for the connections to be established
+				await conn1.getInitializers();
+				await conn2.getInitializers();
+
+				// Get initializers to verify connection params were used
+				const initializers = await conn1.getInitializers();
+
+				// Verify both connection names were recorded
+				expect(initializers).toContain("user1");
+				expect(initializers).toContain("user2");
+
+				// Clean up
+				await conn1.dispose();
+				await conn2.dispose();
+			});
+		});
+
+		describe("Lifecycle Hooks", () => {
+			test("should trigger lifecycle hooks", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Create and connect
+				const connHandle = client.counterWithLifecycle.getOrCreate(
+					["test-lifecycle"],
+					{
+						params: { trackLifecycle: true },
+					},
+				);
+				const connection = connHandle.connect();
+
+				// Verify lifecycle events were triggered
+				const events = await connection.getEvents();
+				expect(events).toEqual(["onStart", "onBeforeConnect", "onConnect"]);
+
+				// Disconnect should trigger onDisconnect
+				await connection.dispose();
+
+				// Reconnect to check if onDisconnect was called
+				const handle = client.counterWithLifecycle.getOrCreate([
+					"test-lifecycle",
+				]);
+				const finalEvents = await handle.getEvents();
+				expect(finalEvents).toBeOneOf([
+					// Still active
+					["onStart", "onBeforeConnect", "onConnect", "onDisconnect"],
+					// Went to sleep and woke back up
+					[
+						"onStart",
+						"onBeforeConnect",
+						"onConnect",
+						"onDisconnect",
+						"onStart",
+					],
+				]);
+			});
+		});
+
+		describe("Connection Liveness", () => {
+			// TODO: KIT-242
+			test.skip("should return correct liveness status for connections", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Create actor and connection
+				const handle = client.connLivenessActor.getOrCreate([
+					"test-liveness-status",
+				]);
+				const connA = handle.connect();
+				const connB = handle.connect();
+
+				const connAId = await connA.getConnectionId();
+				const connBId = await connB.getConnectionId();
+
+				// Verify connection works initially
+				await connA.increment(5);
+				await connB.increment(5);
+
+				const counter = await handle.getCounter();
+				expect(counter).toBe(10);
+
+				const connectionsStatusBeforeKill =
+					await handle.getWsConnectionsLiveness();
+				expect(connectionsStatusBeforeKill).toHaveLength(2);
+				expect(connectionsStatusBeforeKill).toContainEqual(
+					expect.objectContaining({
+						id: connAId,
+						status: "connected",
+						lastSeen: FAKE_TIME.getTime(),
+					}),
+				);
+				expect(connectionsStatusBeforeKill).toContainEqual(
+					expect.objectContaining({
+						id: connBId,
+						status: "connected",
+						lastSeen: FAKE_TIME.getTime(),
+					}),
+				);
+
+				// Kill one connection
+				await handle.kill(connAId); // instead of dispose, we use kill to simulate a disconnection (e.g. network failure)
+				// connA.dispose();
+				// we killed the connection, but the actor instance does not know about it yet
+				// it should still be in the list of connections, but with a status of "reconnecting"
+				const connectionsStatusAfterKill =
+					await handle.getWsConnectionsLiveness();
+				expect(connectionsStatusAfterKill).toEqual(
+					expect.arrayContaining([
+						expect.objectContaining({
+							id: connAId,
+							status: "reconnecting",
+							lastSeen: FAKE_TIME.getTime(),
+						}),
+						expect.objectContaining({
+							id: connBId,
+							status: "connected",
+							lastSeen: FAKE_TIME.getTime(),
+						}),
+					]),
+				);
+
+				// default time to wait for cleanup is 5 seconds
+				// check actor options
+				await waitFor(driverTestConfig, 5_000);
+
+				// After timeout, the killed connection should be unavailable, since the manager has cleaned it up
+				const connectionsStatusAfterCleanup =
+					await handle.getWsConnectionsLiveness();
+				expect(connectionsStatusAfterCleanup).not.toContainEqual(
+					expect.objectContaining({
+						id: connAId,
+					}),
+				);
+				expect(connectionsStatusAfterCleanup).toContainEqual(
+					expect.objectContaining({
+						id: connBId,
+					}),
+				);
+			});
+		});
+	});
+}

package/src/driver-test-suite/tests/actor-driver.ts

@@ -0,0 +1,25 @@
+import { describe } from "vitest";
+import type { DriverTestConfig } from "../mod";
+import { runActorScheduleTests } from "./actor-schedule";
+import { runActorSleepTests } from "./actor-sleep";
+import { runActorStateTests } from "./actor-state";
+
+export function runActorDriverTests(driverTestConfig: DriverTestConfig) {
+	describe("Actor Driver Tests", () => {
+		// Run state persistence tests
+		runActorStateTests(driverTestConfig);
+
+		// Run scheduled alarms tests
+		runActorScheduleTests(driverTestConfig);
+	});
+}
+
+/** Actor driver tests that need to be tested for all transport mechanisms. */
+export function runActorDriverTestsWithTransport(
+	driverTestConfig: DriverTestConfig,
+) {
+	describe("Actor Driver Tests", () => {
+		// Run actor sleep tests
+		runActorSleepTests(driverTestConfig);
+	});
+}