Fix/Relation model class and Self relation class missing imports

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<string>);
+  // 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<string>);
+  }
 
   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<string>);
+  }
+
+  // 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),
             };
           },
         ),

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")
+}

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
+  });
+});

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');
+  });
+});