diff --git a/src/generate-class.ts b/src/generate-class.ts index b66aaf9..9af64db 100644 --- a/src/generate-class.ts +++ b/src/generate-class.ts @@ -221,9 +221,11 @@ function generateRelationClass( } }); - generateRelationImportsImport(sourceFile, [ - ...relationImports, - ] as Array); + if (relationImports.size > 0) { + generateRelationImportsImport(sourceFile, [ + ...relationImports, + ] as Array); + } sourceFile.addClass({ name: `${model.name}Relations`, @@ -264,6 +266,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 +320,7 @@ function generateCombinedClass( hasExclamationToken: field.isRequired, hasQuestionToken: !field.isRequired, trailingTrivia: '\r\n', + decorators: getDecoratorsByFieldType(field, config.swagger), }; }, ), 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'); + }); +});