feat(server): attach .errors into middlewares

unnoq · Feb 9, 2025 · 3783978 · 3783978
1 parent fca153f
commit 3783978
Show file tree

Hide file tree

Showing 5 changed files with 83 additions and 14 deletions.
diff --git a/packages/server/src/builder.test.ts b/packages/server/src/builder.test.ts
@@ -3,11 +3,16 @@ import { baseErrorMap, baseMeta, baseRoute, generalSchema, inputSchema, outputSc
 import { router } from '../tests/shared'
 import { Builder } from './builder'
 import { unlazy } from './lazy'
-import * as MiddlewareDecorated from './middleware-decorated'
+import { decorateMiddleware } from './middleware-decorated'
 import { DecoratedProcedure } from './procedure-decorated'
 import * as Router from './router'
 
-const decorateMiddlewareSpy = vi.spyOn(MiddlewareDecorated, 'decorateMiddleware')
+vi.mock('./middleware-decorated', () => ({
+  decorateMiddleware: vi.fn(mid => ({
+    mapInput: vi.fn(map => [mid, map]),
+  })),
+}))
+
 const adaptRouterSpy = vi.spyOn(Router, 'adaptRouter')
 
 const mid = vi.fn()
@@ -98,9 +103,9 @@ describe('builder', () => {
     const mid = vi.fn()
     const applied = builder.middleware(mid)
 
-    expect(applied).toBe(decorateMiddlewareSpy.mock.results[0]?.value)
-    expect(decorateMiddlewareSpy).toBeCalledTimes(1)
-    expect(decorateMiddlewareSpy).toBeCalledWith(mid)
+    expect(applied).toBe(vi.mocked(decorateMiddleware).mock.results[0]?.value)
+    expect(decorateMiddleware).toBeCalledTimes(1)
+    expect(decorateMiddleware).toBeCalledWith(mid)
   })
 
   it('.errors', () => {
@@ -116,15 +121,51 @@ describe('builder', () => {
     })
   })
 
-  it('.use', () => {
-    const mid2 = vi.fn()
-    const applied = builder.use(mid2)
+  describe('.use', () => {
+    it('without map input', () => {
+      const mid2 = vi.fn()
+      const applied = builder.use(mid2)
+
+      expect(applied).instanceOf(Builder)
+      expect(applied).not.toBe(builder)
+      expect(applied['~orpc']).toEqual({
+        ...def,
+        middlewares: [mid, mid2],
+      })
+    })
 
-    expect(applied).instanceOf(Builder)
-    expect(applied).not.toBe(builder)
-    expect(applied['~orpc']).toEqual({
-      ...def,
-      middlewares: [mid, mid2],
+    it('with map input', () => {
+      const mid2 = vi.fn()
+      const map = vi.fn()
+      const applied = builder.use(mid2, map)
+
+      expect(applied).instanceOf(Builder)
+      expect(applied).not.toBe(builder)
+      expect(applied['~orpc']).toEqual({
+        ...def,
+        middlewares: [mid, [mid2, map]],
+      })
+
+      expect(decorateMiddleware).toBeCalledTimes(1)
+      expect(decorateMiddleware).toBeCalledWith(mid2)
+    })
+
+    it('with attached error map', () => {
+      const errorMap = {
+        ADDITION: {},
+        OVERRIDE: { message: 'this is low priority' },
+      }
+      const mid2 = vi.fn() as any
+      mid2['~attachedErrorMap'] = errorMap
+      const applied = builder.use(mid2)
+
+      expect(applied).instanceOf(Builder)
+      expect(applied).not.toBe(builder)
+      expect(applied['~orpc']).toEqual({
+        ...def,
+        errorMap: { ...errorMap, ...def.errorMap },
+        middlewares: [mid, mid2],
+      })
     })
   })
 

diff --git a/packages/server/src/builder.ts b/packages/server/src/builder.ts
@@ -99,7 +99,9 @@ export class Builder<
 
   middleware<UOutContext extends Context, TInput, TOutput = any>( // = any here is important to make middleware can be used in any output by default
     middleware: Middleware<TCurrentContext, UOutContext, TInput, TOutput, ORPCErrorConstructorMap<TErrorMap>, TMeta>,
-  ): DecoratedMiddleware<TCurrentContext, UOutContext, TInput, TOutput, ORPCErrorConstructorMap<any>, TMeta> { // ORPCErrorConstructorMap<any> ensures middleware can used in any procedure
+  ): DecoratedMiddleware<TCurrentContext, UOutContext, TInput, TOutput, any, TMeta> { // any ensures middleware can used in any procedure without matching error requirements
+    middleware['~attachedErrorMap'] = this['~orpc'].errorMap
+
     return decorateMiddleware(middleware)
   }
 
@@ -147,6 +149,9 @@ export class Builder<
 
     return new Builder({
       ...this['~orpc'],
+      errorMap: mapped['~attachedErrorMap']
+        ? mergeErrorMap(mapped['~attachedErrorMap'], this['~orpc'].errorMap)
+        : this['~orpc'].errorMap,
       middlewares: addMiddleware(this['~orpc'].middlewares, mapped),
     })
   }

diff --git a/packages/server/src/middleware.ts b/packages/server/src/middleware.ts
@@ -46,6 +46,8 @@ export interface Middleware<
   TErrorConstructorMap extends ORPCErrorConstructorMap<any>,
   TMeta extends Meta,
 > {
+  '~attachedErrorMap'?: ErrorMap
+
   (
     options: MiddlewareOptions<TInContext, TOutput, TErrorConstructorMap, TMeta>,
     input: TInput,

diff --git a/packages/server/src/procedure-decorated.test.ts b/packages/server/src/procedure-decorated.test.ts
@@ -96,6 +96,24 @@ describe('decoratedProcedure', () => {
         middlewares: [...def.middlewares, [mid, map]],
       })
     })
+
+    it('with attached error map', () => {
+      const errorMap = {
+        ADDITION: {},
+        OVERRIDE: { message: 'this is low priority' },
+      }
+      const mid2 = vi.fn() as any
+      mid2['~attachedErrorMap'] = errorMap
+      const applied = decorated.use(mid2)
+
+      expect(applied).not.toBe(decorated)
+      expect(applied).toBeInstanceOf(DecoratedProcedure)
+      expect(applied['~orpc']).toEqual({
+        ...def,
+        errorMap: { ...errorMap, ...def.errorMap },
+        middlewares: [...def.middlewares, mid2],
+      })
+    })
   })
 
   it('.callable', () => {

diff --git a/packages/server/src/procedure-decorated.ts b/packages/server/src/procedure-decorated.ts
@@ -84,6 +84,9 @@ export class DecoratedProcedure<
 
     return new DecoratedProcedure({
       ...this['~orpc'],
+      errorMap: mapped['~attachedErrorMap']
+        ? mergeErrorMap(mapped['~attachedErrorMap'], this['~orpc'].errorMap)
+        : this['~orpc'].errorMap,
       middlewares: addMiddleware(this['~orpc'].middlewares, mapped),
     })
   }