impl createPartialContext for tests

__tests__/e2e.test.ts

@@ -1,6 +1,7 @@
 import { assert, beforeEach, describe, expect, test, vi } from 'vitest';
 import {
   closeAllConnections,
+  createPartialContext,
   isReadableDone,
   numberOfConnections,
   readNextResult,
@@ -972,6 +973,72 @@ describe.each(testMatrix())(
       });
     });
 
+    test('createPartialContext throws on unmocked property access', async () => {
+      // eslint-disable-next-line @typescript-eslint/consistent-type-definitions
+      type TestContext = {
+        db: { query: (sql: string) => string };
+        cache: { get: (key: string) => string };
+      };
+
+      const ctx = createPartialContext<TestContext>({
+        db: { query: (sql) => `result: ${sql}` },
+      });
+
+      // provided properties work
+      expect(ctx.db.query('SELECT 1')).toBe('result: SELECT 1');
+
+      // unmocked properties throw
+      expect(() => ctx.cache).toThrow(
+        'cache is not mocked in the test context',
+      );
+    });
+
+    test('createPartialContext works as extendedContext with server dispose', async () => {
+      const clientTransport = getClientTransport('client');
+      const serverTransport = getServerTransport();
+      const dbDispose = vi.fn();
+
+      // eslint-disable-next-line @typescript-eslint/consistent-type-definitions
+      type TestContext = {
+        db: { [Symbol.asyncDispose]: () => Promise<void> };
+        cache: { get: (key: string) => string };
+      };
+
+      const ctx = createPartialContext<TestContext>({
+        db: { [Symbol.asyncDispose]: dbDispose },
+      });
+
+      const ServiceSchema = createServiceSchema<TestContext>();
+      const services = {
+        test: ServiceSchema.define({
+          ping: Procedure.rpc({
+            requestInit: Type.Object({}),
+            responseData: Type.Object({}),
+            async handler() {
+              return Ok({});
+            },
+          }),
+        }),
+      };
+
+      const server = createServer(serverTransport, services, {
+        extendedContext: ctx,
+      });
+      addPostTestCleanup(async () => {
+        await cleanupTransports([clientTransport, serverTransport]);
+      });
+
+      // server.close() should dispose context values without
+      // throwing on unmocked properties (cache)
+      await server.close();
+      expect(dbDispose).toBeCalledTimes(1);
+      await testFinishesCleanly({
+        clientTransports: [clientTransport],
+        serverTransport,
+        server,
+      });
+    });
+
     test('works with non-object schemas', async () => {
       // setup
       const clientTransport = getClientTransport('client');

package.json

@@ -1,7 +1,7 @@
 {
   "name": "@replit/river",
   "description": "It's like tRPC but... with JSON Schema Support, duplex streaming and support for service multiplexing. Transport agnostic!",
-  "version": "0.213.0",
+  "version": "0.213.1",
   "type": "module",
   "exports": {
     ".": {

testUtil/index.ts

@@ -246,3 +246,45 @@ export function closeAllConnections<ConnType extends Connection>(
     conn.close();
   }
 }
+
+/**
+ * Wraps a partial context object in a proxy that throws when accessing
+ * properties that weren't provided. This is useful for test contexts where
+ * you only want to mock the dependencies a test actually uses.
+ *
+ * Symbols and `then` are allowed through without throwing — river checks
+ * for `Symbol.asyncDispose` / `Symbol.dispose` on context values during
+ * `server.close()`, and `then` is checked by the JS runtime when the
+ * proxy is returned from an async function.
+ *
+ * @example
+ * ```ts
+ * const ctx = createPartialContext<MyContext>({
+ *   database: mockDb,
+ *   // accessing ctx.redis will throw
+ * });
+ *
+ * const server = createServer(transport, services, {
+ *   extendedContext: ctx,
+ * });
+ * ```
+ */
+export function createPartialContext<T extends Record<string, unknown>>(
+  partial: Partial<T>,
+): T {
+  return new Proxy(partial as T, {
+    get(target, prop, receiver) {
+      if (prop in target) {
+        return Reflect.get(target, prop, receiver);
+      }
+
+      if (typeof prop === 'string' && prop !== 'then') {
+        throw new Error(
+          `${prop} is not mocked in the test context. Provide it via createPartialContext if your test needs it.`,
+        );
+      }
+
+      return undefined;
+    },
+  });
+}