fix some edge cases

Author: alisabzevari

src/myzod/index.ts

@@ -92,13 +92,17 @@ export const MyZodSchemaVisitor = (schema: GraphQLSchema, config: ValidationSche
     },
     UnionTypeDefinition: (node: UnionTypeDefinitionNode) => {
       const unionName = tsVisitor.convertName(node.name.value);
-      const unionElements = node.types?.map(t => `${t.name.value}Schema()`).join(', ');
+      const unionElements = node.types?.map(t => `${tsVisitor.convertName(t.name.value)}Schema()`).join(', ');
+      const unionElementsCount = node.types?.length ?? 0;
+
+      const union =
+        unionElementsCount > 1 ? indent(`return myzod.union([${unionElements}])`) : indent(`return ${unionElements}`);
 
       const result = new DeclarationBlock({})
         .export()
         .asKind('function')
         .withName(`${unionName}Schema()`)
-        .withBlock(indent(`return myzod.union([${unionElements}])`));
+        .withBlock(union);
 
       return result.string;
     },
@@ -194,6 +198,11 @@ const generateNameNodeMyZodSchema = (
     return `${enumName}Schema`;
   }
 
+  if (typ?.astNode?.kind === 'UnionTypeDefinition') {
+    const enumName = tsVisitor.convertName(typ.astNode.name.value);
+    return `${enumName}Schema()`;
+  }
+
   return myzod4Scalar(config, tsVisitor, node.value);
 };
 

src/zod/index.ts

@@ -103,13 +103,17 @@ export const ZodSchemaVisitor = (schema: GraphQLSchema, config: ValidationSchema
     },
     UnionTypeDefinition: (node: UnionTypeDefinitionNode) => {
       const unionName = tsVisitor.convertName(node.name.value);
-      const unionElements = node.types?.map(t => `${t.name.value}Schema()`).join(', ');
+      const unionElements = node.types?.map(t => `${tsVisitor.convertName(t.name.value)}Schema()`).join(', ');
+      const unionElementsCount = node.types?.length ?? 0;
+
+      const union =
+        unionElementsCount > 1 ? indent(`return z.union([${unionElements}])`) : indent(`return ${unionElements}`);
 
       const result = new DeclarationBlock({})
         .export()
         .asKind('function')
         .withName(`${unionName}Schema()`)
-        .withBlock(indent(`return z.union([${unionElements}])`));
+        .withBlock(union);
 
       return result.string;
     },
@@ -205,6 +209,11 @@ const generateNameNodeZodSchema = (
     return `${enumName}Schema`;
   }
 
+  if (typ?.astNode?.kind === 'UnionTypeDefinition') {
+    const enumName = tsVisitor.convertName(typ.astNode.name.value);
+    return `${enumName}Schema()`;
+  }
+
   return zod4Scalar(config, tsVisitor, node.value);
 };
 

tests/myzod.spec.ts

@@ -570,6 +570,72 @@ describe('myzod', () => {
         expect(result.content).toContain(wantContain);
       }
     });
+
+    it('generate union types with single element', async () => {
+      const schema = buildSchema(/* GraphQL */ `
+        type Square {
+          size: Int
+        }
+        type Circle {
+          radius: Int
+        }
+        union Shape = Circle | Square
+
+        type Geometry {
+          shape: Shape
+        }
+      `);
+
+      const result = await plugin(
+        schema,
+        [],
+        {
+          schema: 'myzod',
+          withObjectType: true,
+        },
+        {}
+      );
+
+      const wantContains = [
+        'export function GeometrySchema(): myzod.Type<Geometry> {',
+        'return myzod.object({',
+        "__typename: myzod.literal('Geometry').optional(),",
+        'shape: ShapeSchema().optional().nullable()',
+        '}',
+      ];
+      for (const wantContain of wantContains) {
+        expect(result.content).toContain(wantContain);
+      }
+    });
+
+    it('correctly reference generated union types', async () => {
+      const schema = buildSchema(/* GraphQL */ `
+        type Circle {
+          radius: Int
+        }
+        union Shape = Circle
+      `);
+
+      const result = await plugin(
+        schema,
+        [],
+        {
+          schema: 'myzod',
+          withObjectType: true,
+        },
+        {}
+      );
+
+      const wantContains = [
+        // Shape Schema
+        'export function ShapeSchema() {',
+        'CircleSchema()',
+        '}',
+      ];
+      for (const wantContain of wantContains) {
+        expect(result.content).toContain(wantContain);
+      }
+    });
   });
 
   it('properly generates custom directive values', async () => {

tests/zod.spec.ts

@@ -668,6 +668,72 @@ describe('zod', () => {
         expect(result.content).toContain(wantContain);
       }
     });
+
+    it('generate union types with single element', async () => {
+      const schema = buildSchema(/* GraphQL */ `
+        type Square {
+          size: Int
+        }
+        type Circle {
+          radius: Int
+        }
+        union Shape = Circle | Square
+
+        type Geometry {
+          shape: Shape
+        }
+      `);
+
+      const result = await plugin(
+        schema,
+        [],
+        {
+          schema: 'zod',
+          withObjectType: true,
+        },
+        {}
+      );
+
+      const wantContains = [
+        'export function GeometrySchema(): z.ZodObject<Properties<Geometry>> {',
+        'return z.object({',
+        "__typename: z.literal('Geometry').optional(),",
+        'shape: ShapeSchema().nullish()',
+        '}',
+      ];
+      for (const wantContain of wantContains) {
+        expect(result.content).toContain(wantContain);
+      }
+    });
+
+    it('correctly reference generated union types', async () => {
+      const schema = buildSchema(/* GraphQL */ `
+        type Circle {
+          radius: Int
+        }
+        union Shape = Circle
+      `);
+
+      const result = await plugin(
+        schema,
+        [],
+        {
+          schema: 'zod',
+          withObjectType: true,
+        },
+        {}
+      );
+
+      const wantContains = [
+        // Shape Schema
+        'export function ShapeSchema() {',
+        'CircleSchema()',
+        '}',
+      ];
+      for (const wantContain of wantContains) {
+        expect(result.content).toContain(wantContain);
+      }
+    });
   });
 
   it('properly generates custom directive values', async () => {