trpc-uwebsockets 11.0.4 → 11.1.2

package/src/server.spec.ts

@@ -1,56 +1,25 @@
-/* eslint-disable @typescript-eslint/no-unused-vars */
-import type {
-  HTTPHeaders,
-  TRPCLink,
-  WebSocketClientOptions,
-} from '@trpc/client';
-import {
-  createTRPCClient,
-  createWSClient,
-  httpBatchLink,
-  TRPCClientError,
-  httpBatchStreamLink,
-  httpSubscriptionLink,
-  wsLink,
-} from '@trpc/client';
+import { TRPCClientError } from '@trpc/client';
 import { initTRPC, TRPCError } from '@trpc/server';
 import { observable } from '@trpc/server/observable';
-import { EventSourcePolyfill } from 'event-source-polyfill';
 import { EventEmitter } from 'events';
-import uWs from 'uWebSockets.js';
-import {
-  afterEach,
-  beforeEach,
-  describe,
-  expect,
-  expectTypeOf,
-  test,
-  vi,
-} from 'vitest';
-import { z } from 'zod';
 
-import {
-  type CreateContextOptions,
-  applyRequestHandler,
-  CreateHandlerOptions,
-} from './requestHandler';
-import { applyWebsocketHandler } from './websockets';
+import { afterEach, describe, expect, expectTypeOf, test, vi } from 'vitest';
+import { z } from 'zod';
+import { testFactory } from './___testHelpers';
 
-const config = {
-  prefix: '/trpc',
-};
+import { type CreateContextOptions } from './requestHandler';
+import { TRPCRequestInfo } from '@trpc/server/unstable-core-do-not-import';
 
 async function createContext({ req, res, info, client }: CreateContextOptions) {
   const user = { name: req.headers.get('username') || 'anonymous' };
 
   await new Promise((resolve) => setTimeout(resolve, 10));
 
-  if (req.headers.has('throw')) {
+  if (info.url?.searchParams.get('throw')) {
     throw new TRPCError({
       code: 'BAD_REQUEST',
-      message: req.headers.get('throw')!,
+      message: info.url!.searchParams.get('throw')!,
     });
-    // throw new Error(req.headers.get('throw')!);
   }
 
   // for websocket context throwing after connection
@@ -61,15 +30,6 @@ async function createContext({ req, res, info, client }: CreateContextOptions) {
     });
   }
 
-  // for websocket context throwing during connection
-  const url = new URL(req.url);
-  if (url.searchParams.has('throw')) {
-    throw new TRPCError({
-      code: 'BAD_REQUEST',
-      message: url.searchParams.get('throw')!,
-    });
-  }
-
   // test writing headers inside the context
   if (info.type === 'query') {
     res.cork(() => {
@@ -79,8 +39,6 @@ async function createContext({ req, res, info, client }: CreateContextOptions) {
   return { req, res, user, info, client };
 }
 
-type Context = Awaited<ReturnType<typeof createContext>>;
-
 interface Message {
   id: string;
 }
@@ -90,7 +48,7 @@ function createAppRouter() {
   const onNewMessageSubscription = vi.fn();
   const onSubscriptionEnded = vi.fn();
 
-  const t = initTRPC.context<Context>().create();
+  const t = initTRPC.create();
   const router = t.router;
   const publicProcedure = t.procedure;
 
@@ -110,7 +68,7 @@ function createAppRouter() {
           .nullish()
       )
       .query(({ input, ctx }) => ({
-        text: `hello ${input?.username ?? ctx.user?.name ?? 'world'}`,
+        text: `hello ${input?.username ?? (ctx as any).user?.name ?? 'world'}`,
       })),
     helloMutation: publicProcedure
       .input(z.string())
@@ -126,7 +84,7 @@ function createAppRouter() {
         })
       )
       .mutation(async ({ input, ctx }) => {
-        if (ctx.user.name === 'anonymous') {
+        if ((ctx as any).user.name === 'anonymous') {
           return { error: 'Unauthorized user' };
         }
         const { id, data } = input;
@@ -148,10 +106,12 @@ function createAppRouter() {
       return sub;
     }),
     pubSub: publicProcedure.input(z.null()).subscription(({ ctx }) => {
-      if (!ctx.client) {
+      // TODO: fix all these
+      const anyContext = ctx as any;
+      if (!anyContext.client) {
         throw new Error('assertion: client should never be null in websockets');
       }
-      const client = ctx.client!;
+      const client = anyContext.client!;
       // for some reason client.publish does not work...
       // currently okay is false, not sure why uWs behaves this way
       // publishing from the server works though
@@ -186,7 +146,7 @@ function createAppRouter() {
     }),
     request: router({
       info: publicProcedure.query(({ ctx }) => {
-        return ctx.info;
+        return (ctx as any).info as TRPCRequestInfo;
       }),
     }),
     deferred: publicProcedure
@@ -222,7 +182,10 @@ function createAppRouter() {
             return;
           }
           if (input.throw === i) {
-            throw new Error(i.toString());
+            throw new TRPCError({
+              code: 'BAD_REQUEST',
+              message: i.toString(),
+            });
           }
           yield i;
         }
@@ -232,249 +195,35 @@ function createAppRouter() {
   return { appRouter, ee, onNewMessageSubscription, onSubscriptionEnded };
 }
 
-type CreateAppRouter = Awaited<ReturnType<typeof createAppRouter>>;
-type AppRouter = CreateAppRouter['appRouter'];
-
-interface ServerOptions {
-  appRouter: AppRouter;
-}
-
-function createServer(opts: ServerOptions) {
-  const instance = uWs.App();
-
-  const router = opts.appRouter;
-
-  applyRequestHandler(instance, {
-    prefix: config.prefix,
-    ssl: false,
-    trpcOptions: {
-      router,
-      createContext,
-      onError(data) {
-        // console.error('trpc error', data);
-      },
-      responseMeta({ ctx, paths, type, errors }) {
-        return {
-          headers: {
-            'Access-Control-Allow-Origin': '*',
-          },
-        };
-      },
-    } satisfies CreateHandlerOptions<AppRouter>['trpcOptions'],
-  });
-  applyWebsocketHandler(instance, {
-    prefix: config.prefix,
-    router,
-    createContext,
-  });
-
-  instance.get('/hello', async (res, req) => {
-    res.end('Hello world');
-  });
-  instance.post('/hello', async (res, req) => {
-    res.end(JSON.stringify({ hello: 'POST', body: 'TODO, why?' }));
-  });
+let stopFn = () => {};
 
-  instance.ws('/ws', {
-    message: (client, rawMsg) => {
-      client.send(rawMsg);
-    },
-  });
-
-  instance.ws('/pubsub', {
-    open: (client) => {
-      const ok = client.subscribe('topic');
-
-      if (!ok) {
-        throw new Error("assertion: failed to subscribe to 'topic'");
-      }
-    },
-  });
-
-  instance.any('/*', (res, req) => {
-    res.writeStatus('404 NOT FOUND');
-    res.end();
-  });
+afterEach(() => {
+  stopFn();
+});
 
-  let socket: uWs.us_listen_socket | false | null = null;
+function defaultFactory(config?: {
+  createContext?: (opts: CreateContextOptions) => Promise<any>;
+}) {
+  const router = createAppRouter();
 
-  instance.listen('0.0.0.0', 0, (token) => {
-    socket = token;
+  const factoryVal = testFactory({
+    appRouter: router.appRouter,
+    createContext: config?.createContext ?? createContext,
   });
 
-  if (!socket) {
-    throw new Error('could not make a socket');
-  }
-
-  const port = uWs.us_socket_local_port(socket);
+  stopFn = factoryVal.stop;
 
   return {
-    stop() {
-      if (!socket) {
-        throw new Error('could not close socket as socket is already closed');
-      }
-      uWs.us_listen_socket_close(socket);
-      socket = null;
-    },
-    port,
-    instance,
+    ...router,
+    ...factoryVal,
   };
 }
 
-const orderedResults: number[] = [];
-const linkSpy: TRPCLink<AppRouter> = () => {
-  // here we just got initialized in the app - this happens once per app
-  // useful for storing cache for instance
-  return ({ next, op }) => {
-    // this is when passing the result to the next link
-    // each link needs to return an observable which propagates results
-    return observable((observer) => {
-      const unsubscribe = next(op).subscribe({
-        next(value) {
-          orderedResults.push(value.result.data as number);
-          observer.next(value);
-        },
-        error: observer.error,
-      });
-      return unsubscribe;
-    });
-  };
-};
-
-interface ClientOptions {
-  wsClientOptions?: Omit<WebSocketClientOptions, 'url'> | undefined;
-  queryParams?: Record<string, string> | undefined;
-  headers?: HTTPHeaders | undefined;
-  port: number;
-}
-
-function toQueryString(queryParams: Record<string, string>) {
-  const raw = new URLSearchParams(queryParams).toString();
-  return raw.length == 0 ? '' : `?${raw}`;
-}
-
-function createClientWs(opts: ClientOptions) {
-  const qs = toQueryString(opts.queryParams ?? {});
-  const host = `localhost:${opts.port}${config.prefix}${qs}`;
-  const wsClient = createWSClient({
-    retryDelayMs: () => 99999, // never retry by default
-    ...opts.wsClientOptions,
-    url: `ws://${host}`,
-  });
-  const client = createTRPCClient<AppRouter>({
-    links: [linkSpy, wsLink({ client: wsClient })],
-  });
-
-  return { client, wsClient };
-}
-
-function createClientBatchStream(opts: ClientOptions) {
-  const qs = toQueryString(opts.queryParams ?? {});
-  const host = `localhost:${opts.port}${config.prefix}${qs}`;
-  const client = createTRPCClient<AppRouter>({
-    links: [
-      linkSpy,
-      httpBatchStreamLink({
-        url: `http://${host}`,
-        headers: opts.headers,
-      }),
-    ],
-  });
-
-  return { client };
-}
-
-function createClientBatch(opts: ClientOptions) {
-  const qs = toQueryString(opts.queryParams ?? {});
-  const host = `localhost:${opts.port}${config.prefix}${qs}`;
-  const client = createTRPCClient<AppRouter>({
-    links: [
-      httpBatchLink({
-        url: `http://${host}`,
-        headers: opts.headers,
-      }),
-    ],
-  });
-
-  return { client };
-}
-
-function createClientSSE(opts: ClientOptions) {
-  const qs = toQueryString(opts.queryParams ?? {});
-  const host = `localhost:${opts.port}${config.prefix}${qs}`;
-  const client = createTRPCClient<AppRouter>({
-    links: [
-      // loggerLink(),
-      httpSubscriptionLink({
-        url: `http://${host}`,
-        connectionParams: opts.wsClientOptions?.connectionParams,
-        // ponyfill EventSource
-        EventSource: EventSourcePolyfill as any,
-      }),
-    ],
-  });
-  return { client };
-}
-type ClientType = 'batchStream' | 'batch' | 'sse';
-async function createApp(serverOptions?: Partial<ServerOptions> | undefined) {
-  const { appRouter, ee } = createAppRouter();
-  const { instance, port, stop } = createServer({
-    ...(serverOptions ?? {}),
-    appRouter,
-  });
-
-  return {
-    server: instance,
-    stop,
-    getClient(
-      clientType: ClientType,
-      clientOptions?: Partial<ClientOptions> | undefined
-    ) {
-      switch (clientType) {
-        case 'batchStream':
-          return createClientBatchStream({
-            ...clientOptions,
-            port,
-          });
-        case 'batch':
-          return createClientBatch({
-            ...clientOptions,
-            port,
-          });
-        case 'sse':
-          return createClientSSE({
-            ...clientOptions,
-            port,
-          });
-        default:
-          throw new Error('unknown client');
-      }
-    },
-    getClientWs(clientOptions?: Partial<ClientOptions> | undefined) {
-      return createClientWs({
-        ...clientOptions,
-        port,
-      });
-    },
-    ee,
-    port,
-    opts: serverOptions,
-  };
-}
-
-let app: Awaited<ReturnType<typeof createApp>>;
-beforeEach(async () => {
-  orderedResults.length = 0;
-  app = await createApp();
-});
-
-afterEach(() => {
-  app.stop();
-});
-
 describe('server', () => {
   test('fetch GET smoke', async () => {
-    const res = await fetch(`http://localhost:${app.port}/hello`, {
+    const ctx = defaultFactory();
+
+    const res = await ctx.fetch(`/hello`, {
       method: 'GET',
       headers: {
         Accept: 'application/json',
@@ -485,39 +234,23 @@ describe('server', () => {
     expect(await res.text()).toEqual('Hello world');
   });
 
-  test('response meta', async () => {
-    const res = await fetch(
-      `http://localhost:${app.port}/trpc/ping?input=${encodeURI('{}')}`
-    );
+  test('forwards response meta', async () => {
+    const ctx = defaultFactory();
+    const res = await ctx.fetch(`/trpc/ping?input=${encodeURI('{}')}`, {});
     expect(res.status).toEqual(200);
     expect(res.headers.get('Access-Control-Allow-Origin')).toEqual('*'); // from the meta
     expect(res.headers.get('x-test')).toEqual('true'); // from the context
   });
 
   test('query', async () => {
-    {
-      // batch stream
-      const { client } = app.getClient('batchStream');
-      expect(await client.ping.query()).toMatchInlineSnapshot(`"pong"`);
-      expect(await client.hello.query()).toMatchInlineSnapshot(`
-          Object {
-            "text": "hello anonymous",
-          }
-      `);
-      expect(
-        await client.hello.query({
-          username: 'test',
-        })
-      ).toMatchInlineSnapshot(`
-          Object {
-            "text": "hello test",
-          }
-      `);
-    }
-
-    {
-      // batch
-      const { client } = app.getClient('batch');
+    const ctx = defaultFactory();
+    const clients = [
+      ctx.clientBatchStream().client,
+      ctx.clientBatch().client,
+      ctx.clientWs().client,
+    ];
+
+    for (const client of clients) {
       expect(await client.ping.query()).toMatchInlineSnapshot(`"pong"`);
       expect(await client.hello.query()).toMatchInlineSnapshot(`
           Object {
@@ -537,23 +270,14 @@ describe('server', () => {
   });
 
   test('mutation', async () => {
-    {
-      // batch stream
-      const { client } = app.getClient('batchStream');
-      expect(
-        await client.editPost.mutate({
-          id: '42',
-          data: { title: 'new_title', text: 'new_text' },
-        })
-      ).toMatchInlineSnapshot(`
-      Object {
-        "error": "Unauthorized user",
-      }
-    `);
-    }
-    {
-      // batch
-      const { client } = app.getClient('batch');
+    const ctx = defaultFactory();
+    const clients = [
+      ctx.clientBatchStream().client,
+      ctx.clientBatch().client,
+      ctx.clientWs().client,
+    ];
+
+    for (const client of clients) {
       expect(
         await client.editPost.mutate({
           id: '42',
@@ -567,33 +291,102 @@ describe('server', () => {
     }
   });
 
-  test('streaming', async () => {
-    const { client } = app.getClient('batchStream');
-    const results = await Promise.all([
-      client.deferred.query({ wait: 3 }),
-      client.deferred.query({ wait: 1 }),
-      client.deferred.query({ wait: 2 }),
-    ]);
-    expect(results).toEqual([3, 1, 2]);
-    expect(orderedResults).toEqual([1, 2, 3]);
-  });
+  test('subscription', { timeout: 5000 }, async () => {
+    const ctx = defaultFactory();
+    const ee = ctx.ee;
+    const clients = [ctx.clientSse().client, ctx.clientWs().client];
 
-  test('batched requests in body work correctly', async () => {
-    {
-      // batch stream
-      const { client } = app.getClient('batchStream');
+    for (const client of clients) {
+      ee.once('subscription:created', () => {
+        setTimeout(() => {
+          ee.emit('server:msg', {
+            id: '1',
+          });
+          ee.emit('server:msg', {
+            id: '2',
+          });
+        });
+      });
+      const onStartedMock = vi.fn();
+      const onDataMock = vi.fn();
+      const sub = client.onMessage.subscribe('onMessage', {
+        onStarted: onStartedMock,
+        onData(data) {
+          expectTypeOf(data).not.toBeAny();
+          expectTypeOf(data).toMatchTypeOf<Message>();
+          onDataMock(data);
+        },
+      });
 
-      const res = await Promise.all([
-        client.helloMutation.mutate('world'),
-        client.helloMutation.mutate('KATT'),
+      await vi.waitFor(
+        () => {
+          expect(onStartedMock).toHaveBeenCalledTimes(1);
+          expect(onDataMock).toHaveBeenCalledTimes(2);
+        },
+        { timeout: 3000 }
+      );
+
+      ee.emit('server:msg', {
+        id: '3',
+      });
+
+      await vi.waitFor(() => {
+        expect(onDataMock).toHaveBeenCalledTimes(3);
+      });
+
+      expect(onDataMock.mock.calls).toMatchInlineSnapshot(`
+        Array [
+          Array [
+            Object {
+              "id": "1",
+            },
+          ],
+          Array [
+            Object {
+              "id": "2",
+            },
+          ],
+          Array [
+            Object {
+              "id": "3",
+            },
+          ],
+        ]
+      `);
+
+      sub.unsubscribe();
+
+      await vi.waitFor(() => {
+        expect(ee.listenerCount('server:msg')).toBe(0);
+        expect(ee.listenerCount('server:error')).toBe(0);
+      });
+    }
+  });
+
+  test('streaming', async () => {
+    const ctx = defaultFactory();
+    const tests = [ctx.clientBatchStream(), ctx.clientWs()];
+
+    for (const { client, orderedResults } of tests) {
+      const results = await Promise.all([
+        client.deferred.query({ wait: 3 }),
+        client.deferred.query({ wait: 1 }),
+        client.deferred.query({ wait: 2 }),
       ]);
-      expect(res).toEqual(['hello world', 'hello KATT']);
+      expect(results).toEqual([3, 1, 2]);
+      expect(orderedResults).toEqual([1, 2, 3]);
     }
+  });
 
-    {
-      //batch
-      const { client } = app.getClient('batch');
+  test('batched requests', async () => {
+    const ctx = defaultFactory();
+    const clients = [
+      ctx.clientBatchStream().client,
+      ctx.clientBatch().client,
+      ctx.clientWs().client,
+    ];
 
+    for (const client of clients) {
       const res = await Promise.all([
         client.helloMutation.mutate('world'),
         client.helloMutation.mutate('KATT'),
@@ -602,105 +395,84 @@ describe('server', () => {
     }
   });
 
-  test('handles throwing procedure', async () => {
-    // batch stream
-    const { client } = app.getClient('batchStream');
-    await expect(
-      client.throw.query('expected_procedure_error')
-    ).rejects.toThrowErrorMatchingInlineSnapshot(
-      `[TRPCClientError: expected_procedure_error]`
-    );
-  });
-
   test('handles throwing context', async () => {
-    const { client } = app.getClient('batchStream', {
-      headers: {
+    const ctx = defaultFactory();
+    const clientOpts = {
+      queryParams: {
         throw: 'expected_context_error',
       },
-    });
+    };
+    const clients = [
+      ctx.clientBatchStream(clientOpts).client,
+      ctx.clientBatch(clientOpts).client,
+    ];
+
+    for (const client of clients) {
+      await expect(
+        client.echo.query('hii')
+      ).rejects.toThrowErrorMatchingInlineSnapshot(
+        `[TRPCClientError: expected_context_error]`
+      );
+    }
+
+    // websocket context resolution failure hides the message
+    const wsClient = ctx.clientWs(clientOpts).client;
     await expect(
-      client.echo.query('hii')
+      wsClient.echo.query('hii')
     ).rejects.toThrowErrorMatchingInlineSnapshot(
-      `[TRPCClientError: expected_context_error]`
+      `[TRPCClientError: Unknown error]`
     );
   });
 
-  // TODO: test failure of context as in v10
-
-  test('subscription - sse', { timeout: 5000 }, async () => {
-    const { client } = app.getClient('sse');
-
-    app.ee.once('subscription:created', () => {
-      setTimeout(() => {
-        app.ee.emit('server:msg', {
-          id: '1',
-        });
-        app.ee.emit('server:msg', {
-          id: '2',
-        });
-      });
-    });
-
-    const onStartedMock = vi.fn();
-    const onDataMock = vi.fn();
-    const sub = client.onMessage.subscribe('onMessage', {
-      onStarted: onStartedMock,
-      onData(data) {
-        expectTypeOf(data).not.toBeAny();
-        expectTypeOf(data).toMatchTypeOf<Message>();
-        onDataMock(data);
-      },
-    });
-
-    await vi.waitFor(
-      () => {
-        expect(onStartedMock).toHaveBeenCalledTimes(1);
-        expect(onDataMock).toHaveBeenCalledTimes(2);
-      },
-      { timeout: 3000 }
-    );
-
-    app.ee.emit('server:msg', {
-      id: '3',
-    });
-
-    await vi.waitFor(() => {
-      expect(onDataMock).toHaveBeenCalledTimes(3);
-    });
+  test('handles throwing procedure', async () => {
+    const ctx = defaultFactory();
+    const clients = [
+      ctx.clientBatchStream().client,
+      ctx.clientBatch().client,
+      ctx.clientWs().client,
+    ];
+
+    for (const client of clients) {
+      await expect(
+        client.throw.query('expected_procedure_error')
+      ).rejects.toThrowErrorMatchingInlineSnapshot(
+        `[TRPCClientError: expected_procedure_error]`
+      );
+      expect(ctx.onServerErrorSpy).toHaveBeenCalled();
+    }
+  });
 
-    expect(onDataMock.mock.calls).toMatchInlineSnapshot(`
-      Array [
-        Array [
-          Object {
-            "id": "1",
-          },
-        ],
-        Array [
-          Object {
-            "id": "2",
-          },
-        ],
-        Array [
-          Object {
-            "id": "3",
+  test('handles throwing subscription', async () => {
+    const ctx = defaultFactory();
+    const clients = [ctx.clientWs(), ctx.clientSse()];
+
+    for (const { client, orderedResults } of clients) {
+      let error: TRPCClientError<any> | null = null;
+      client.count.subscribe(
+        {
+          from: 0,
+          count: 10,
+          throw: 5,
+        },
+        {
+          onError(err) {
+            error = err;
           },
-        ],
-      ]
-    `);
-
-    sub.unsubscribe();
-
-    await vi.waitFor(() => {
-      expect(app.ee.listenerCount('server:msg')).toBe(0);
-      expect(app.ee.listenerCount('server:error')).toBe(0);
-    });
+        }
+      );
+      await vi.waitFor(() => {
+        expect(orderedResults).toEqual([0, 1, 2, 3, 4]);
+        expect(ctx.onServerErrorSpy).toHaveBeenCalled();
+        expect(error!.message).toEqual(new TRPCClientError('5').message);
+      });
+    }
   });
 });
 
 describe('websocket', () => {
   test('does not bind other websocket connection', async () => {
-    const client = new WebSocket(`ws://localhost:${app.port}/ws`);
-
+    const ctx = defaultFactory();
+    const client = new WebSocket(`ws://localhost:${ctx.port}/ws`);
     await new Promise<void>((resolve, reject) => {
       client.onopen = () => {
         client.send('hello');
@@ -708,113 +480,47 @@ describe('websocket', () => {
       };
       client.onerror = reject;
     });
-
     const promise = new Promise<string>((resolve) => {
       client.onmessage = (msg) => {
         return resolve(msg.data);
       };
     });
-
     const message = await promise;
-
     expect(message.toString()).toBe('hello');
-
     client.close();
   });
 
-  test('basic functionality', async () => {
-    const { client } = app.getClientWs();
+  test('properly passes context', async () => {
+    const contextDone = vi.fn();
+    const ctx = defaultFactory({
+      createContext: async ({ req, res, info, client }) => {
+        expect(client != null);
+        expect(res != null);
+        expect(client).equals(res);
 
-    app.ee.once('subscription:created', () => {
-      setTimeout(() => {
-        app.ee.emit('server:msg', {
-          id: '1',
-        });
-        app.ee.emit('server:msg', {
-          id: '2',
-        });
-      });
-    });
-
-    const onStartedMock = vi.fn();
-    const onDataMock = vi.fn();
-    const sub = client.onMessage.subscribe('onMessage', {
-      onStarted: onStartedMock,
-      onData(data) {
-        expectTypeOf(data).not.toBeAny();
-        expectTypeOf(data).toMatchTypeOf<Message>();
-        onDataMock(data);
-      },
-    });
+        // can still get the removeAddress from res
+        // hopefully this stays constant in ci?
+        const ip = new TextDecoder().decode(res.getRemoteAddressAsText());
+        expect(ip).equals('127.0.0.1');
 
-    await vi.waitFor(() => {
-      expect(onStartedMock).toHaveBeenCalledTimes(1);
-      expect(onDataMock).toHaveBeenCalledTimes(2);
-    });
+        // but other res-like features are not possible
+        expect('writeHeader' in res).equals(false);
 
-    app.ee.emit('server:msg', {
-      id: '3',
-    });
+        contextDone();
 
-    await vi.waitFor(() => {
-      expect(onDataMock).toHaveBeenCalledTimes(3);
-    });
-
-    expect(onDataMock.mock.calls).toMatchInlineSnapshot(`
-      Array [
-        Array [
-          Object {
-            "id": "1",
-          },
-        ],
-        Array [
-          Object {
-            "id": "2",
-          },
-        ],
-        Array [
-          Object {
-            "id": "3",
-          },
-        ],
-      ]
-    `);
-
-    sub.unsubscribe();
-
-    await vi.waitFor(() => {
-      expect(app.ee.listenerCount('server:msg')).toBe(0);
-      expect(app.ee.listenerCount('server:error')).toBe(0);
-    });
-  });
-
-  test('handles throwing procedure', async () => {
-    const { client } = app.getClientWs();
-
-    let error: any = null;
-
-    client.count.subscribe(
-      {
-        from: 0,
-        count: 10,
-        throw: 5,
+        return { req, res, info, client };
       },
-      {
-        onError(err) {
-          error = err;
-        },
-      }
-    );
-
-    await vi.waitFor(() => {
-      expect(error.message).toEqual(new TRPCClientError('5').message);
     });
+    const client = ctx.clientWs().client;
+    await client.echo.query('testing');
+
+    expect(contextDone).toHaveBeenCalledOnce();
   });
 
   test('handles throwing context with connection params', async () => {
+    const server = defaultFactory();
     let closeCode: number | undefined = undefined;
-
-    const { client } = app.getClientWs({
+    const { client } = server.clientWs({
       wsClientOptions: {
         connectionParams: {
           throw: 'expected_context_error',
@@ -828,28 +534,20 @@ describe('websocket', () => {
       },
     });
 
-    try {
-      await client.echo.query('hi');
-    } catch (err) {
-      //
-    }
-    // client.count.subscribe(
-    //   {
-    //     from: 0,
-    //     count: 10,
-    //   },
-    //   {}
-    // );
+    await expect(
+      client.echo.query('hi')
+    ).rejects.toThrowErrorMatchingInlineSnapshot(
+      `[TRPCClientError: Unknown error]`
+    );
 
     await vi.waitFor(() => {
       expect(closeCode).toEqual(1008);
     });
   });
-
   test('handles throwing context without connection params', async () => {
+    const server = defaultFactory();
     let closeCode: number | undefined = undefined;
-
-    const { client } = app.getClientWs({
+    const { client } = server.clientWs({
       queryParams: {
         throw: 'expected_context_error',
       },
@@ -860,12 +558,10 @@ describe('websocket', () => {
       },
     });
 
-    client.count.subscribe(
-      {
-        from: 0,
-        count: 10,
-      },
-      {}
+    await expect(
+      client.echo.query('hi')
+    ).rejects.toThrowErrorMatchingInlineSnapshot(
+      `[TRPCClientError: Unknown error]`
     );
 
     await vi.waitFor(() => {
@@ -873,78 +569,31 @@ describe('websocket', () => {
     });
   });
 
-  // for some reason client.publish does not work, even though it should
-  test('uWebsockets pubsub', { skip: true }, async () => {
-    const clientWs = new WebSocket(`ws://localhost:${app.port}/pubsub`);
-
-    const messages: string[] = [];
-    let connected = false;
-    clientWs.onmessage = (ev) => {
-      console.log('triggered onmessage', ev.data);
-      messages.push(JSON.parse(ev.data));
-    };
-    clientWs.onerror = () => {
-      throw new Error('client had an error');
-    };
-
-    clientWs.onopen = () => {
-      connected = true;
-    };
-
-    await vi.waitFor(() => {
-      expect(connected).toBe(true);
-    });
-
-    expect(app.server.numSubscribers('topic')).toBe(1);
-
-    const { client } = app.getClientWs();
-    const sub = client.pubSub.subscribe(null, {
-      onData(data) {
-        console.log('trpc onMessage', data);
-
-        expectTypeOf(data).not.toBeAny();
-        expectTypeOf(data).toMatchTypeOf<Message>();
-        messages.push(data.id);
-      },
-    });
-
-    await vi.waitFor(() => {
-      expect(messages).toEqual(['created', 'message', 'ended']);
-    });
-
-    sub.unsubscribe();
-    clientWs.close();
-  });
-
   test('keep alive', async () => {
+    const ctx = defaultFactory();
     let closeCode: number | undefined = undefined;
-
-    const { client, wsClient } = app.getClientWs({
+    const { client, wsClient } = ctx.clientWs({
       wsClientOptions: {
         onClose(cause) {
           closeCode = cause?.code;
         },
         keepAlive: {
           enabled: true,
-          intervalMs: 200,
-          pongTimeoutMs: 400,
+          intervalMs: 100,
+          pongTimeoutMs: 200,
         },
       },
     });
-
     const { unsubscribe } = client.waitForMs.subscribe(2_000, {});
-
     await vi.waitFor(() => {
       expect(wsClient.connection!.state).toBe('open');
     });
-
     let pongCount = 0;
     wsClient.connection!.ws!.addEventListener('message', (ev) => {
       if (ev.data === 'PONG') {
         pongCount++;
       }
     });
-
     await vi.waitFor(
       () => {
         expect(pongCount).toEqual(4);
@@ -954,13 +603,44 @@ describe('websocket', () => {
         timeout: 2_000,
       }
     );
-
     unsubscribe();
-
     wsClient.close();
-
     await vi.waitFor(() => {
       expect(closeCode).toEqual(1005);
     });
   });
+  //   // for some reason client.publish does not work, even though it should
+  //   test('uWebsockets pubsub', { skip: true }, async () => {
+  //     const clientWs = new WebSocket(`ws://localhost:${app.port}/pubsub`);
+  //     const messages: string[] = [];
+  //     let connected = false;
+  //     clientWs.onmessage = (ev) => {
+  //       console.log('triggered onmessage', ev.data);
+  //       messages.push(JSON.parse(ev.data));
+  //     };
+  //     clientWs.onerror = () => {
+  //       throw new Error('client had an error');
+  //     };
+  //     clientWs.onopen = () => {
+  //       connected = true;
+  //     };
+  //     await vi.waitFor(() => {
+  //       expect(connected).toBe(true);
+  //     });
+  //     expect(app.server.numSubscribers('topic')).toBe(1);
+  //     const { client } = app.clientWs();
+  //     const sub = client.pubSub.subscribe(null, {
+  //       onData(data) {
+  //         console.log('trpc onMessage', data);
+  //         expectTypeOf(data).not.toBeAny();
+  //         expectTypeOf(data).toMatchTypeOf<Message>();
+  //         messages.push(data.id);
+  //       },
+  //     });
+  //     await vi.waitFor(() => {
+  //       expect(messages).toEqual(['created', 'message', 'ended']);
+  //     });
+  //     sub.unsubscribe();
+  //     clientWs.close();
+  //   });
 });