Add flag to register entities created

src/factory.ts

@@ -1,4 +1,4 @@
-import type { DataSource, SaveOptions } from 'typeorm'
+import type { DataSource, RemoveOptions, SaveOptions } from 'typeorm'
 import { EagerInstanceAttribute, LazyInstanceAttribute } from './instanceAttributes'
 import { BaseSubfactory } from './subfactories'
 import type { Constructable, FactorizedAttrs } from './types'
@@ -8,14 +8,16 @@ export abstract class Factory<T> {
   protected abstract dataSource: DataSource
   protected abstract attrs(): FactorizedAttrs<T>
 
+  private createdEntities: T[] = []
+
   /**
    * Make a new entity without persisting it
    */
   async make(overrideParams: Partial<FactorizedAttrs<T>> = {}): Promise<T> {
     const attrs = { ...this.attrs(), ...overrideParams }
 
-    const entity = await this.makeEntity(attrs, false)
-    await this.applyLazyAttributes(entity, attrs, false)
+    const entity = await this.makeEntity(attrs, false, false)
+    await this.applyLazyAttributes(entity, attrs, false, false)
 
     return entity
   }
@@ -31,72 +33,156 @@ export abstract class Factory<T> {
     return list
   }
 
+  public async create(overrideParams?: Partial<FactorizedAttrs<T>>): Promise<T>
+  public async create(overrideParams?: Partial<FactorizedAttrs<T>>, shouldRegister?: boolean): Promise<T>
+  public async create(overrideParams?: Partial<FactorizedAttrs<T>>, saveOptions?: SaveOptions): Promise<T>
+  public async create(
+    overrideParams?: Partial<FactorizedAttrs<T>>,
+    saveOptions?: SaveOptions,
+    shouldRegister?: boolean,
+  ): Promise<T>
+
   /**
    * Create a new entity and persist it
    */
-  async create(overrideParams: Partial<FactorizedAttrs<T>> = {}, saveOptions?: SaveOptions): Promise<T> {
+  public async create(
+    overrideParams?: Partial<FactorizedAttrs<T>>,
+    saveOptionsOrShouldRegister?: SaveOptions | boolean,
+    shouldRegister?: boolean,
+  ): Promise<T> {
+    const saveOptions = typeof saveOptionsOrShouldRegister === 'object' ? saveOptionsOrShouldRegister : undefined
+    const shouldRegisterComputed =
+      typeof saveOptionsOrShouldRegister === 'boolean' ? saveOptionsOrShouldRegister : shouldRegister || false
+
     const attrs = { ...this.attrs(), ...overrideParams }
     const preloadedAttrs = Object.entries(attrs).filter(([, value]) => !(value instanceof LazyInstanceAttribute))
 
-    const entity = await this.makeEntity(Object.fromEntries(preloadedAttrs) as FactorizedAttrs<T>, true)
+    const entity = await this.makeEntity(
+      Object.fromEntries(preloadedAttrs) as FactorizedAttrs<T>,
+      true,
+      shouldRegisterComputed,
+    )
 
     const em = this.dataSource.createEntityManager()
     const savedEntity = await em.save<T>(entity, saveOptions)
+    if (shouldRegisterComputed) this.createdEntities.push(savedEntity)
+
+    await this.applyLazyAttributes(savedEntity, attrs, true, shouldRegisterComputed)
 
-    await this.applyLazyAttributes(savedEntity, attrs, true)
     return em.save<T>(savedEntity, saveOptions)
   }
 
+  public async createMany(amount: number): Promise<T[]>
+  public async createMany(amount: number, overrideParams?: Partial<FactorizedAttrs<T>>): Promise<T[]>
+  public async createMany(
+    amount: number,
+    overrideParams?: Partial<FactorizedAttrs<T>>,
+    shouldRegister?: boolean,
+  ): Promise<T[]>
+  public async createMany(
+    amount: number,
+    overrideParams?: Partial<FactorizedAttrs<T>>,
+    saveOptions?: SaveOptions,
+  ): Promise<T[]>
+  public async createMany(
+    amount: number,
+    overrideParams?: Partial<FactorizedAttrs<T>>,
+    saveOptions?: SaveOptions,
+    shouldRegister?: boolean,
+  ): Promise<T[]>
+
   /**
    * Create many new entities and persist them
    */
   async createMany(
     amount: number,
     overrideParams: Partial<FactorizedAttrs<T>> = {},
-    saveOptions?: SaveOptions,
+    saveOptionsOrShouldRegister?: SaveOptions | boolean,
+    shouldRegister?: boolean,
   ): Promise<T[]> {
+    const saveOptions = typeof saveOptionsOrShouldRegister === 'object' ? saveOptionsOrShouldRegister : undefined
+    const shouldRegisterComputed =
+      typeof saveOptionsOrShouldRegister === 'boolean' ? saveOptionsOrShouldRegister : shouldRegister || false
+
     const list = []
     for (let index = 0; index < amount; index++) {
-      list[index] = await this.create(overrideParams, saveOptions)
+      list[index] = await this.create(overrideParams, saveOptions, shouldRegisterComputed)
     }
     return list
   }
 
-  private async makeEntity(attrs: FactorizedAttrs<T>, shouldPersist: boolean) {
+  public getCreatedEntities() {
+    return this.createdEntities
+  }
+
+  /**
+   * This method deletes all entities that were created by this factory.
+   * The order of deletion is the reverse of the order of creation.
+   * @experimental As of version 1.2.0
+   */
+  public async cleanUp(removeOptions?: RemoveOptions) {
+    const entityManager = this.dataSource.createEntityManager()
+    for (const entity of this.createdEntities.reverse()) {
+      await entityManager.remove(entity, removeOptions)
+    }
+  }
+
+  public flushEntities() {
+    this.createdEntities = []
+  }
+
+  private async makeEntity(attrs: FactorizedAttrs<T>, shouldPersist: boolean, shouldRegister: boolean) {
     const entity = new this.entity()
 
     await Promise.all(
-      Object.entries(attrs).map(async ([key, value]) => {
-        Object.assign(entity, { [key]: await Factory.resolveValue(value, shouldPersist) })
-      }),
+      Object.entries(attrs)
+        .filter(([, value]) => {
+          return !(value instanceof EagerInstanceAttribute || value instanceof LazyInstanceAttribute)
+        })
+        .map(async ([key, value]) => {
+          Object.assign(entity, { [key]: await this.resolveValue(value, shouldPersist, shouldRegister) })
+        }),
     )
 
     await Promise.all(
-      Object.entries(attrs).map(async ([key, value]) => {
-        if (value instanceof EagerInstanceAttribute) {
-          const newAttrib = value.resolve(entity)
-          Object.assign(entity, { [key]: await Factory.resolveValue(newAttrib, shouldPersist) })
-        }
-      }),
+      Object.entries(attrs)
+        .filter(([, value]) => value instanceof EagerInstanceAttribute)
+        .map(async ([key, value]) => {
+          if (value instanceof EagerInstanceAttribute) {
+            const newAttrib = value.resolve(entity)
+            Object.assign(entity, { [key]: await this.resolveValue(newAttrib, shouldPersist, shouldRegister) })
+          }
+        }),
     )
 
     return entity
   }
 
-  private async applyLazyAttributes(entity: T, attrs: FactorizedAttrs<T>, shouldPersist: boolean) {
+  private async applyLazyAttributes(
+    entity: T,
+    attrs: FactorizedAttrs<T>,
+    shouldPersist: boolean,
+    shouldRegister: boolean,
+  ) {
     await Promise.all(
-      Object.entries(attrs).map(async ([key, value]) => {
-        if (value instanceof LazyInstanceAttribute) {
-          const newAttrib = value.resolve(entity)
-          Object.assign(entity, { [key]: await Factory.resolveValue(newAttrib, shouldPersist) })
-        }
-      }),
+      Object.entries(attrs)
+        .filter(([, value]) => value instanceof LazyInstanceAttribute)
+        .map(async ([key, value]) => {
+          if (value instanceof LazyInstanceAttribute) {
+            const newAttrib = value.resolve(entity)
+            Object.assign(entity, { [key]: await this.resolveValue(newAttrib, shouldPersist, shouldRegister) })
+          }
+        }),
     )
   }
 
-  private static resolveValue(value: unknown, shouldPersist: boolean) {
+  private async resolveValue(value: unknown, shouldPersist: boolean, shouldRegister: boolean) {
     if (value instanceof BaseSubfactory) {
-      return shouldPersist ? value.create() : value.make()
+      if (!shouldPersist) return value.make()
+
+      const [entity, createdEntities] = await value.createAndFlush(shouldRegister)
+      if (shouldRegister) this.createdEntities.push(...createdEntities)
+      return entity
     } else if (value instanceof Function) {
       return value()
     } else {

src/index.ts

@@ -1,4 +1,4 @@
 export * from './factory'
 export * from './instanceAttributes'
-export * from './subfactories'
+export { CollectionSubfactory, SingleSubfactory } from './subfactories'
 export * from './types'

src/subfactories/baseSubfactory.ts

@@ -8,6 +8,14 @@ export abstract class BaseSubfactory<T> {
     this.factoryInstance = new factory()
   }
 
-  abstract create(): Promise<T> | Promise<T[]>
+  public async createAndFlush(shouldRegister?: boolean) {
+    const entity = await this.create(shouldRegister)
+    const result = [entity, this.factoryInstance.getCreatedEntities()]
+    this.factoryInstance.flushEntities()
+
+    return result
+  }
+
+  abstract create(shouldRegister?: boolean): Promise<T> | Promise<T[]>
   abstract make(): Promise<T> | Promise<T[]>
 }

src/subfactories/collectionSubfactory.ts

@@ -7,8 +7,8 @@ export class CollectionSubfactory<T> extends BaseSubfactory<T> {
     super(factory, values)
   }
 
-  create() {
-    return this.factoryInstance.createMany(this.count, this.values)
+  create(shouldRegister?: boolean) {
+    return this.factoryInstance.createMany(this.count, this.values, shouldRegister)
   }
 
   make() {

src/subfactories/singleSubfactory.ts

@@ -7,8 +7,8 @@ export class SingleSubfactory<T> extends BaseSubfactory<T> {
     super(factory, values)
   }
 
-  create() {
-    return this.factoryInstance.create(this.values)
+  create(shouldRegister?: boolean) {
+    return this.factoryInstance.create(this.values, shouldRegister)
   }
 
   make() {

test/factory.test.ts

@@ -90,6 +90,12 @@ describe(Factory, () => {
 
         expect(userMaked1).not.toStrictEqual(userMaked2)
       })
+
+      test('Should not register maked entity', async () => {
+        await factory.make()
+
+        expect(factory.getCreatedEntities()).toHaveLength(0)
+      })
     })
 
     describe(PetFactory, () => {
@@ -134,6 +140,11 @@ describe(Factory, () => {
     describe(UserFactory, () => {
       const factory = new UserFactory()
 
+      beforeEach(() => {
+        factory.cleanUp()
+        factory.flushEntities()
+      })
+
       test('Should create a new entity', async () => {
         const userCreated = await factory.create()
 
@@ -215,6 +226,63 @@ describe(Factory, () => {
 
         expect(userCreated1).not.toStrictEqual(userCreated2)
       })
+
+      test('Should not fail if save options are provided', async () => {
+        const userCreated = await factory.create(
+          {
+            name: 'john',
+          },
+          {
+            reload: true,
+          },
+        )
+
+        expect(userCreated.name).toBe('john')
+      })
+
+      test('Should not register created entity', async () => {
+        await factory.create()
+
+        expect(factory.getCreatedEntities()).toHaveLength(0)
+      })
+
+      test('Should register created entity', async () => {
+        await factory.create({}, true)
+
+        expect(factory.getCreatedEntities()).toHaveLength(1)
+      })
+
+      test('Should register created entity with subfactory', async () => {
+        await factory.create(
+          {
+            pets: new LazyInstanceAttribute((instance) => new CollectionSubfactory(PetFactory, 1, { owner: instance })),
+          },
+          true,
+        )
+
+        expect(factory.getCreatedEntities()).toHaveLength(2)
+      })
+
+      test('Should remove created entity with helper function', async () => {
+        await factory.create(
+          {
+            pets: new LazyInstanceAttribute((instance) => new CollectionSubfactory(PetFactory, 1, { owner: instance })),
+          },
+          true,
+        )
+
+        const em = dataSource.createEntityManager()
+
+        const totalUsers = await em.count(User)
+        const totalPets = await em.count(Pet)
+
+        await factory.cleanUp()
+
+        const totalUsersAfterCleanUp = await em.count(User)
+        expect(totalUsersAfterCleanUp).toBe(totalUsers - 1)
+        const totalPetsAfterCleanUp = await em.count(Pet)
+        expect(totalPetsAfterCleanUp).toBe(totalPets - 1)
+      })
     })
 
     describe(PetFactory, () => {
@@ -253,5 +321,16 @@ describe(Factory, () => {
         expect(entity.id).toBeDefined()
       })
     })
+
+    test('Should not fail if save options are provided', async () => {
+      const count = 2
+      const factory = new UserFactory()
+      const entitiesCreated = await factory.createMany(count, {}, { reload: true })
+
+      expect(entitiesCreated).toHaveLength(count)
+      entitiesCreated.forEach((entity) => {
+        expect(entity.id).toBeDefined()
+      })
+    })
   })
 })