fix: Fix wrap-cache was not working (#29)

* fix: Fix wrap-cache was not working * fix: Use original function as cache key
toss · Nov 16, 2023 · b3145a6 · b3145a6
1 parent 72ad78b
commit b3145a6
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 3 deletions.
diff --git a/src/__test__/aop.module.test.ts b/src/__test__/aop.module.test.ts
@@ -1,6 +1,6 @@
 import 'reflect-metadata';
 
-import { Controller, Get, Injectable, Module } from '@nestjs/common';
+import { Controller, Get, Injectable, Module, Scope } from '@nestjs/common';
 import { FastifyAdapter } from '@nestjs/platform-fastify';
 import { Test } from '@nestjs/testing';
 import { AopModule } from '../aop.module';
@@ -150,6 +150,50 @@ describe('AopModule', () => {
     expect(fooService.multipleDecorated()).toMatchInlineSnapshot(`"012"`);
   });
 
+  it('Wrap should be executed only once in default scope', async () => {
+    let wrapped = 0;
+    @Injectable({ scope: Scope.DEFAULT })
+    class FooService {
+      @AopTesting({
+        wrapCallback: () => {
+          wrapped++;
+        },
+      })
+      decorated() {
+        return '0';
+      }
+    }
+
+    @Module({
+      providers: [FooService],
+      exports: [FooService],
+    })
+    class FooModule {}
+
+    const module = await Test.createTestingModule({
+      imports: [
+        AopModule,
+        FooModule,
+        AopTestingModule.registerAsync({
+          imports: [FooModule],
+          inject: [FooService],
+          useFactory: (fooService: FooService) => {
+            return [fooService];
+          },
+        }),
+      ],
+    }).compile();
+
+    const app = module.createNestApplication(new FastifyAdapter());
+    await app.init();
+    const fooService = app.get(FooService);
+    fooService.decorated();
+    fooService.decorated();
+    fooService.decorated();
+
+    expect(wrapped).toBe(1);
+  });
+
   /**
    * There are codes that using `function.name`.
    * Therefore the codes below are necessary.

diff --git a/src/auto-aspect-executor.ts b/src/auto-aspect-executor.ts
@@ -86,7 +86,8 @@ export class AutoAspectExecutor implements OnModuleInit {
     // eslint-disable-next-line @typescript-eslint/no-this-alias
     const self = this;
     const wrappedFn = function (this: object, ...args: unknown[]) {
-      const cached = self.wrappedMethodCache.get(aopMetadata);
+      const cache = self.wrappedMethodCache.get(this) || new WeakMap();
+      const cached = cache.get(originalFn);
       if (cached) {
         return cached.apply(this, args);
       }
@@ -97,7 +98,8 @@ export class AutoAspectExecutor implements OnModuleInit {
         method: originalFn.bind(this),
         metadata,
       });
-      self.wrappedMethodCache.set(this, wrappedMethod);
+      cache.set(originalFn, wrappedMethod);
+      self.wrappedMethodCache.set(this, cache);
       return wrappedMethod.apply(this, args);
     };