rivetkit 2.0.2 → 2.0.4

package/src/driver-test-suite/mod.ts

@@ -0,0 +1,239 @@
+import { serve as honoServe } from "@hono/node-server";
+import { createNodeWebSocket, type NodeWebSocket } from "@hono/node-ws";
+import { bundleRequire } from "bundle-require";
+import invariant from "invariant";
+import { describe } from "vitest";
+import type { Transport } from "@/client/mod";
+import { createManagerRouter } from "@/manager/router";
+import type { DriverConfig, Registry, RunConfig } from "@/mod";
+import { RunConfigSchema } from "@/registry/run-config";
+import { getPort } from "@/test/mod";
+import { logger } from "./log";
+import { runActionFeaturesTests } from "./tests/action-features";
+import { runActorConnTests } from "./tests/actor-conn";
+import { runActorConnStateTests } from "./tests/actor-conn-state";
+import {
+	runActorDriverTests,
+	runActorDriverTestsWithTransport,
+} from "./tests/actor-driver";
+import { runActorErrorHandlingTests } from "./tests/actor-error-handling";
+import { runActorHandleTests } from "./tests/actor-handle";
+import { runActorInlineClientTests } from "./tests/actor-inline-client";
+import { runActorInspectorTests } from "./tests/actor-inspector";
+import { runActorMetadataTests } from "./tests/actor-metadata";
+import { runActorOnStateChangeTests } from "./tests/actor-onstatechange";
+import { runActorVarsTests } from "./tests/actor-vars";
+import { runManagerDriverTests } from "./tests/manager-driver";
+import { runRawHttpTests } from "./tests/raw-http";
+import { runRawHttpRequestPropertiesTests } from "./tests/raw-http-request-properties";
+import { runRawWebSocketTests } from "./tests/raw-websocket";
+import { runRequestAccessTests } from "./tests/request-access";
+
+export interface SkipTests {
+	schedule?: boolean;
+	sleep?: boolean;
+}
+
+export interface DriverTestConfig {
+	/** Deploys an registry and returns the connection endpoint. */
+	start(projectDir: string): Promise<DriverDeployOutput>;
+
+	/**
+	 * If we're testing with an external system, we should use real timers
+	 * instead of Vitest's mocked timers.
+	 **/
+	useRealTimers?: boolean;
+
+	/** Cloudflare Workers has some bugs with cleanup. */
+	HACK_skipCleanupNet?: boolean;
+
+	skip?: SkipTests;
+
+	transport?: Transport;
+
+	clientType: ClientType;
+
+	cleanup?: () => Promise<void>;
+}
+
+/**
+ * The type of client to run the test with.
+ *
+ * The logic for HTTP vs inline is very different, so this helps validate all behavior matches.
+ **/
+type ClientType = "http" | "inline";
+
+export interface DriverDeployOutput {
+	endpoint: string;
+	namespace: string;
+	runnerName: string;
+
+	/** Cleans up the test. */
+	cleanup(): Promise<void>;
+}
+
+/** Runs all Vitest tests against the provided drivers. */
+export function runDriverTests(
+	driverTestConfigPartial: Omit<DriverTestConfig, "clientType" | "transport">,
+) {
+	for (const clientType of ["http", "inline"] as ClientType[]) {
+		const driverTestConfig: DriverTestConfig = {
+			...driverTestConfigPartial,
+			clientType,
+		};
+
+		describe(`client type (${clientType})`, () => {
+			runActorDriverTests(driverTestConfig);
+			runManagerDriverTests(driverTestConfig);
+
+			for (const transport of ["websocket", "sse"] as Transport[]) {
+				describe(`transport (${transport})`, () => {
+					runActorConnTests({
+						...driverTestConfig,
+						transport,
+					});
+
+					runActorConnStateTests({ ...driverTestConfig, transport });
+
+					runRequestAccessTests({ ...driverTestConfig, transport });
+
+					runActorDriverTestsWithTransport({ ...driverTestConfig, transport });
+				});
+			}
+
+			runActorHandleTests(driverTestConfig);
+
+			runActionFeaturesTests(driverTestConfig);
+
+			runActorVarsTests(driverTestConfig);
+
+			runActorMetadataTests(driverTestConfig);
+
+			runActorOnStateChangeTests(driverTestConfig);
+
+			runActorErrorHandlingTests(driverTestConfig);
+
+			runActorInlineClientTests(driverTestConfig);
+
+			runRawHttpTests(driverTestConfig);
+
+			runRawHttpRequestPropertiesTests(driverTestConfig);
+
+			runRawWebSocketTests(driverTestConfig);
+
+			// TODO: re-expose this once we can have actor queries on the gateway
+			// runRawHttpDirectRegistryTests(driverTestConfig);
+
+			// TODO: re-expose this once we can have actor queries on the gateway
+			// runRawWebSocketDirectRegistryTests(driverTestConfig);
+
+			runActorInspectorTests(driverTestConfig);
+		});
+	}
+}
+
+/**
+ * Helper function to adapt the drivers to the Node.js runtime for tests.
+ *
+ * This is helpful for drivers that run in-process as opposed to drivers that rely on external tools.
+ */
+export async function createTestRuntime(
+	registryPath: string,
+	driverFactory: (registry: Registry<any>) => Promise<{
+		rivetEngine?: { endpoint: string; namespace: string; runnerName: string };
+		driver: DriverConfig;
+		cleanup?: () => Promise<void>;
+	}>,
+): Promise<DriverDeployOutput> {
+	const {
+		mod: { registry },
+	} = await bundleRequire<{ registry: Registry<any> }>({
+		filepath: registryPath,
+	});
+
+	// TODO: Find a cleaner way of flagging an registry as test mode (ideally not in the config itself)
+	// Force enable test
+	registry.config.test.enabled = true;
+
+	// Build drivers
+	const {
+		driver,
+		cleanup: driverCleanup,
+		rivetEngine,
+	} = await driverFactory(registry);
+
+	if (rivetEngine) {
+		// TODO: We don't need createTestRuntime fort his
+		// Using external Rivet engine
+
+		const cleanup = async () => {
+			await driverCleanup?.();
+		};
+
+		return {
+			endpoint: rivetEngine.endpoint,
+			namespace: rivetEngine.namespace,
+			runnerName: rivetEngine.runnerName,
+			cleanup,
+		};
+	} else {
+		// Start server for Rivet engine
+
+		// Build driver config
+		// biome-ignore lint/style/useConst: Assigned later
+		let upgradeWebSocket: any;
+		const config: RunConfig = RunConfigSchema.parse({
+			driver,
+			getUpgradeWebSocket: () => upgradeWebSocket!,
+			inspector: {
+				enabled: true,
+				token: () => "token",
+			},
+		});
+
+		// Create router
+		const managerDriver = driver.manager(registry.config, config);
+		const { router } = createManagerRouter(
+			registry.config,
+			config,
+			managerDriver,
+			false,
+		);
+
+		// Inject WebSocket
+		const nodeWebSocket = createNodeWebSocket({ app: router });
+		upgradeWebSocket = nodeWebSocket.upgradeWebSocket;
+
+		// Start server
+		const port = await getPort();
+		const server = honoServe({
+			fetch: router.fetch,
+			hostname: "127.0.0.1",
+			port,
+		});
+		invariant(
+			nodeWebSocket.injectWebSocket !== undefined,
+			"should have injectWebSocket",
+		);
+		nodeWebSocket.injectWebSocket(server);
+		const serverEndpoint = `http://127.0.0.1:${port}`;
+
+		logger().info({ msg: "test serer listening", port });
+
+		// Cleanup
+		const cleanup = async () => {
+			// Stop server
+			await new Promise((resolve) => server.close(() => resolve(undefined)));
+
+			// Extra cleanup
+			await driverCleanup?.();
+		};
+
+		return {
+			endpoint: serverEndpoint,
+			namespace: "default",
+			runnerName: "rivetkit",
+			cleanup,
+		};
+	}
+}

package/src/driver-test-suite/tests/action-features.ts

@@ -0,0 +1,136 @@
+import { describe, expect, test } from "vitest";
+import type { ActorError } from "@/client/errors";
+import type { DriverTestConfig } from "../mod";
+import { setupDriverTest } from "../utils";
+
+export function runActionFeaturesTests(driverTestConfig: DriverTestConfig) {
+	describe("Action Features", () => {
+		// TODO: These do not work with fake timers
+		describe("Action Timeouts", () => {
+			const usesFakeTimers = !driverTestConfig.useRealTimers;
+
+			test("should timeout actions that exceed the configured timeout", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// The quick action should complete successfully
+				const quickResult = await client.shortTimeoutActor
+					.getOrCreate()
+					.quickAction();
+				expect(quickResult).toBe("quick response");
+
+				// The slow action should throw a timeout error
+				await expect(
+					client.shortTimeoutActor.getOrCreate().slowAction(),
+				).rejects.toThrow("Action timed out");
+			});
+
+			test("should respect the default timeout", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// This action should complete within the default timeout
+				const result = await client.defaultTimeoutActor
+					.getOrCreate()
+					.normalAction();
+				expect(result).toBe("normal response");
+			});
+
+			test("non-promise action results should not be affected by timeout", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// Synchronous action should not be affected by timeout
+				const result = await client.syncTimeoutActor.getOrCreate().syncAction();
+				expect(result).toBe("sync response");
+			});
+
+			test("should allow configuring different timeouts for different actors", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				// The short timeout actor should fail
+				await expect(
+					client.shortTimeoutActor.getOrCreate().slowAction(),
+				).rejects.toThrow("Action timed out");
+
+				// The longer timeout actor should succeed
+				const result = await client.longTimeoutActor
+					.getOrCreate()
+					.delayedAction();
+				expect(result).toBe("delayed response");
+			});
+		});
+
+		describe("Action Sync & Async", () => {
+			test("should support synchronous actions", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				const instance = client.syncActionActor.getOrCreate();
+
+				// Test increment action
+				let result = await instance.increment(5);
+				expect(result).toBe(5);
+
+				result = await instance.increment(3);
+				expect(result).toBe(8);
+
+				// Test getInfo action
+				const info = await instance.getInfo();
+				expect(info.currentValue).toBe(8);
+				expect(typeof info.timestamp).toBe("number");
+
+				// Test reset action (void return)
+				await instance.reset();
+				result = await instance.increment(0);
+				expect(result).toBe(0);
+			});
+
+			test("should support asynchronous actions", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				const instance = client.asyncActionActor.getOrCreate();
+
+				// Test delayed increment
+				const result = await instance.delayedIncrement(5);
+				expect(result).toBe(5);
+
+				// Test fetch data
+				const data = await instance.fetchData("test-123");
+				expect(data.id).toBe("test-123");
+				expect(typeof data.timestamp).toBe("number");
+
+				// Test successful async operation
+				const success = await instance.asyncWithError(false);
+				expect(success).toBe("Success");
+
+				// Test error in async operation
+				try {
+					await instance.asyncWithError(true);
+					expect.fail("did not error");
+				} catch (error) {
+					expect((error as ActorError).message).toBe("Intentional error");
+				}
+			});
+
+			test("should handle promises returned from actions correctly", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				const instance = client.promiseActor.getOrCreate();
+
+				// Test resolved promise
+				const resolvedValue = await instance.resolvedPromise();
+				expect(resolvedValue).toBe("resolved value");
+
+				// Test delayed promise
+				const delayedValue = await instance.delayedPromise();
+				expect(delayedValue).toBe("delayed value");
+
+				// Test rejected promise
+				await expect(instance.rejectedPromise()).rejects.toThrow(
+					"promised rejection",
+				);
+
+				// Check state was updated by the delayed promise
+				const results = await instance.getResults();
+				expect(results).toContain("delayed");
+			});
+		});
+	});
+}

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();
+			});
+		});
+	});
+}