diff --git a/src/generate-class.ts b/src/generate-class.ts index b66aaf9..decb6b1 100644 --- a/src/generate-class.ts +++ b/src/generate-class.ts @@ -215,15 +215,31 @@ function generateRelationClass( } const relationImports = new Set(); + let hasSelfRelation = false; + relationFields.forEach((field) => { - if (field.relationName && model.name !== field.type) { - relationImports.add(field.type); + if (field.relationName) { + if (model.name !== field.type) { + relationImports.add(field.type); + } else { + hasSelfRelation = true; + } } }); - generateRelationImportsImport(sourceFile, [ - ...relationImports, - ] as Array); + // For self-relations in the Relations class, import the combined model class + if (hasSelfRelation) { + sourceFile.addImportDeclaration({ + moduleSpecifier: `./${model.name}.model`, + namedImports: [model.name], + }); + } + + if (relationImports.size > 0) { + generateRelationImportsImport(sourceFile, [ + ...relationImports, + ] as Array); + } sourceFile.addClass({ name: `${model.name}Relations`, @@ -264,6 +280,31 @@ function generateCombinedClass( namedImports: [`${model.name}Base`], }); + // Add class validator imports for relation fields + const validatorImports = [ + ...new Set( + relationFields + .map((field) => getDecoratorsImportsByType(field)) + .flatMap((item) => item), + ), + ]; + + if (validatorImports.length > 0) { + generateClassValidatorImport(sourceFile, validatorImports as Array); + } + + // Add Swagger imports if enabled + if ( + config.swagger && + relationFields.length > 0 && + shouldImportSwagger(relationFields as PrismaDMMF.Field[]) + ) { + const swaggerImports = getSwaggerImportsByType( + relationFields as PrismaDMMF.Field[], + ); + generateSwaggerImport(sourceFile, swaggerImports); + } + // Import relation types for the combined class const relationImports = new Set(); relationFields.forEach((field) => { @@ -293,6 +334,7 @@ function generateCombinedClass( hasExclamationToken: field.isRequired, hasQuestionToken: !field.isRequired, trailingTrivia: '\r\n', + decorators: getDecoratorsByFieldType(field, config.swagger), }; }, ), diff --git a/tests/schemas/self-relation.prisma b/tests/schemas/self-relation.prisma new file mode 100644 index 0000000..bdb1da9 --- /dev/null +++ b/tests/schemas/self-relation.prisma @@ -0,0 +1,27 @@ +generator client { + provider = "prisma-client-js" +} + +generator class_validator { + provider = "node ./lib/generator.js" + output = "../generated/self-relation" + swagger = "true" + separateRelationFields = "true" +} + +datasource db { + provider = "sqlite" + url = "file:./test.db" +} + +model User { + id Int @id @default(autoincrement()) + email String @unique + name String? + createdAt DateTime @default(now()) + + // Self-relation: User can have a mentor (one-to-many) + mentorId Int? + mentor User? @relation("UserMentor", fields: [mentorId], references: [id]) + mentees User[] @relation("UserMentor") +} \ No newline at end of file diff --git a/tests/self-relation.test.ts b/tests/self-relation.test.ts new file mode 100644 index 0000000..67631a5 --- /dev/null +++ b/tests/self-relation.test.ts @@ -0,0 +1,130 @@ +import { exec } from 'child_process'; +import { promisify } from 'util'; +import { existsSync, readFileSync } from 'fs'; +import { describe, it, expect, beforeAll } from 'vitest'; +import path from 'path'; + +const execAsync = promisify(exec); + +describe('Self-Relation Generation', () => { + beforeAll(async () => { + // Build the generator first + await execAsync('npm run build'); + + // Generate models for self-relation schema + const schemaPath = path.resolve(__dirname, 'schemas/self-relation.prisma'); + await execAsync(`npx prisma generate --schema="${schemaPath}"`); + }, 60000); + + it('should generate UserBase class without self-relations', () => { + const outputPath = path.resolve(__dirname, 'generated/self-relation'); + const userBasePath = path.join(outputPath, 'models', 'UserBase.model.ts'); + + expect(existsSync(userBasePath)).toBe(true); + const userBase = readFileSync(userBasePath, 'utf-8'); + + // Should have scalar fields + expect(userBase).toContain('id!: number'); + expect(userBase).toContain('email!: string'); + expect(userBase).toContain('name?: string'); + expect(userBase).toContain('createdAt!: Date'); + expect(userBase).toContain('mentorId?: number'); + + // Should have decorators + expect(userBase).toContain('@IsInt()'); + expect(userBase).toContain('@IsString()'); + expect(userBase).toContain('@IsDate()'); + expect(userBase).toContain('@IsDefined()'); + expect(userBase).toContain('@IsOptional()'); + + // Should NOT have self-relation fields + expect(userBase).not.toContain('mentor?'); + expect(userBase).not.toContain('mentees'); + expect(userBase).not.toContain('User[]'); + }); + + it('should generate UserRelations class with only self-relation fields', () => { + const outputPath = path.resolve(__dirname, 'generated/self-relation'); + const userRelationsPath = path.join( + outputPath, + 'models', + 'UserRelations.model.ts', + ); + + expect(existsSync(userRelationsPath)).toBe(true); + const userRelations = readFileSync(userRelationsPath, 'utf-8'); + + // Should have self-relation fields + expect(userRelations).toContain('mentor?: User'); + expect(userRelations).toContain('mentees!: User[]'); + + // Should have class-validator decorators + expect(userRelations).toContain('@IsOptional()'); + expect(userRelations).toContain('@IsDefined()'); + + // Should have Swagger decorators + expect(userRelations).toContain('@ApiProperty({'); + expect(userRelations).toContain('type: () => User'); + expect(userRelations).toContain('required: false'); + expect(userRelations).toContain('isArray: true'); + + // For self-relations in Relations class, should import from the combined model + expect(userRelations).toContain('import { User } from "./User.model"'); + + // Should NOT have scalar fields + expect(userRelations).not.toContain('id!: number'); + expect(userRelations).not.toContain('email!: string'); + expect(userRelations).not.toContain('mentorId'); + }); + + it('should generate combined User class with self-relations', () => { + const outputPath = path.resolve(__dirname, 'generated/self-relation'); + const userModelPath = path.join(outputPath, 'models', 'User.model.ts'); + + expect(existsSync(userModelPath)).toBe(true); + const userModel = readFileSync(userModelPath, 'utf-8'); + + // Should import from UserBase + expect(userModel).toContain('import { UserBase } from "./UserBase.model"'); + + // Should extend UserBase + expect(userModel).toContain('export class User extends UserBase'); + + // Should have self-relation fields with decorators + expect(userModel).toContain('mentor?: User'); + expect(userModel).toContain('mentees!: User[]'); + + // Should have class-validator imports + expect(userModel).toContain( + 'import { IsOptional, IsDefined } from "class-validator"', + ); + + // Should have Swagger imports + expect(userModel).toContain( + 'import { ApiProperty } from "@nestjs/swagger"', + ); + + // Should have decorators on relation fields + expect(userModel).toContain('@IsOptional()'); + expect(userModel).toContain('@IsDefined()'); + expect(userModel).toContain('@ApiProperty({'); + + // Should NOT import User from itself (no circular import) + expect(userModel).not.toContain('import { User } from "./"'); + expect(userModel).not.toContain('import { User } from "./User.model"'); + }); + + it('should handle self-relations without circular import issues', () => { + const outputPath = path.resolve(__dirname, 'generated/self-relation'); + + // Check that the index file exports User + const indexPath = path.join(outputPath, 'models', 'index.ts'); + expect(existsSync(indexPath)).toBe(true); + const index = readFileSync(indexPath, 'utf-8'); + + // When separateRelationFields is enabled, index should export all classes + expect(index).toContain('export { User } from "./User.model"'); + // Note: UserBase and UserRelations might not be in index if only User is exported + // This depends on the generator's index generation logic + }); +}); diff --git a/tests/swagger-generation.test.ts b/tests/swagger-generation.test.ts index b1c9ec3..97c91d5 100644 --- a/tests/swagger-generation.test.ts +++ b/tests/swagger-generation.test.ts @@ -70,3 +70,124 @@ describe('Swagger Generation', () => { expect(userModel).toContain('@ApiProperty({'); }); }); + +describe('Combined Class with separateRelationFields', () => { + beforeAll(async () => { + // Build the generator first + await execAsync('npm run build'); + + // Generate models for full-features schema (has both swagger and separateRelationFields) + const schemaPath = path.resolve(__dirname, 'schemas/full-features.prisma'); + await execAsync(`npx prisma generate --schema="${schemaPath}"`); + }, 60000); + + it('should generate combined class with class-validator and Swagger decorators for relations', () => { + const outputPath = path.resolve(__dirname, 'generated/full-features'); + + // Check User combined class + const userModelPath = path.join(outputPath, 'models', 'User.model.ts'); + const userModel = readFileSync(userModelPath, 'utf-8'); + + // Check that it imports from UserBase + expect(userModel).toContain('import { UserBase } from "./UserBase.model"'); + + // Check for class-validator imports + expect(userModel).toContain('import { IsDefined } from "class-validator"'); + + // Check for Swagger imports + expect(userModel).toContain( + 'import { ApiProperty } from "@nestjs/swagger"', + ); + + // Check for relation type imports (from index file) + expect(userModel).toContain('import { Post } from "./"'); + + // Check that class extends UserBase + expect(userModel).toContain('export class User extends UserBase'); + + // Check for relation field with decorators + expect(userModel).toContain('posts!: Post[]'); + expect(userModel).toContain('@IsDefined()'); + expect(userModel).toContain( + '@ApiProperty({ isArray: true, type: () => Post })', + ); + }); + + it('should generate Post combined class with decorators for author relation', () => { + const outputPath = path.resolve(__dirname, 'generated/full-features'); + + // Check Post combined class + const postModelPath = path.join(outputPath, 'models', 'Post.model.ts'); + const postModel = readFileSync(postModelPath, 'utf-8'); + + // Check that it imports from PostBase + expect(postModel).toContain('import { PostBase } from "./PostBase.model"'); + + // Check for class-validator imports for optional relation + expect(postModel).toContain('import { IsOptional } from "class-validator"'); + + // Check for Swagger imports + expect(postModel).toContain( + 'import { ApiProperty } from "@nestjs/swagger"', + ); + + // Check for relation type imports (from index file) + expect(postModel).toContain('import { User } from "./"'); + + // Check that class extends PostBase + expect(postModel).toContain('export class Post extends PostBase'); + + // Check for optional relation field with decorators + expect(postModel).toContain('author?: User'); + expect(postModel).toContain('@IsOptional()'); + expect(postModel).toContain( + '@ApiProperty({ required: false, type: () => User })', + ); + }); + + it('should generate base classes without relation fields', () => { + const outputPath = path.resolve(__dirname, 'generated/full-features'); + + // Check UserBase class + const userBasePath = path.join(outputPath, 'models', 'UserBase.model.ts'); + const userBase = readFileSync(userBasePath, 'utf-8'); + + // Should have scalar fields with decorators + expect(userBase).toContain('@IsInt()'); + expect(userBase).toContain('@IsString()'); + expect(userBase).toContain('@ApiProperty({'); + expect(userBase).toContain('id!: number'); + expect(userBase).toContain('email!: string'); + expect(userBase).toContain('name?: string'); + + // Should NOT have relation fields + expect(userBase).not.toContain('posts'); + expect(userBase).not.toContain('Post[]'); + }); + + it('should generate relation classes with only relation fields', () => { + const outputPath = path.resolve(__dirname, 'generated/full-features'); + + // Check UserRelations class + const userRelationsPath = path.join( + outputPath, + 'models', + 'UserRelations.model.ts', + ); + + // Check if file exists + expect(existsSync(userRelationsPath)).toBe(true); + + const userRelations = readFileSync(userRelationsPath, 'utf-8'); + + // Should have relation field with decorators + expect(userRelations).toContain('export class UserRelations'); + expect(userRelations).toContain('posts!: Post[]'); + expect(userRelations).toContain('@IsDefined()'); + expect(userRelations).toContain('@ApiProperty({'); + + // Should NOT have scalar fields + expect(userRelations).not.toContain('id!: number'); + expect(userRelations).not.toContain('email!: string'); + }); +});