Add support for union in myzod

Co-authored-by: alexandruluca <[email protected]>

Author: alisabzevari

src/myzod/index.ts

@@ -9,6 +9,7 @@ import {
   ObjectTypeDefinitionNode,
   EnumTypeDefinitionNode,
   FieldDefinitionNode,
+  UnionTypeDefinitionNode,
 } from 'graphql';
 import { DeclarationBlock, indent } from '@graphql-codegen/visitor-plugin-common';
 import { TsVisitor } from '@graphql-codegen/typescript';
@@ -89,6 +90,18 @@ export const MyZodSchemaVisitor = (schema: GraphQLSchema, config: ValidationSche
         .withName(`${enumname}Schema`)
         .withContent(`myzod.enum(${enumname})`).string;
     },
+    UnionTypeDefinition: (node: UnionTypeDefinitionNode) => {
+      const unionName = tsVisitor.convertName(node.name.value);
+      const unionElements = node.types?.map(t => `${t.name.value}Schema()`).join(', ');
+
+      const result = new DeclarationBlock({})
+        .export()
+        .asKind('function')
+        .withName(`${unionName}Schema()`)
+        .withBlock(indent(`return myzod.union([${unionElements}])`));
+
+      return result.string;
+    },
   };
 };
 

tests/myzod.spec.ts

@@ -538,6 +538,38 @@ describe('myzod', () => {
         expect(result.content).not.toContain(wantNotContain);
       }
     });
+
+    it('generate union types', async () => {
+      const schema = buildSchema(/* GraphQL */ `
+        type Square {
+          size: Int
+        }
+        type Circle {
+          radius: Int
+        }
+        union Shape = Circle | Square
+      `);
+
+      const result = await plugin(
+        schema,
+        [],
+        {
+          schema: 'myzod',
+          withObjectType: true,
+        },
+        {}
+      );
+
+      const wantContains = [
+        // Shape Schema
+        'export function ShapeSchema() {',
+        'myzod.union([CircleSchema(), SquareSchema()])',
+        '}',
+      ];
+      for (const wantContain of wantContains) {
+        expect(result.content).toContain(wantContain);
+      }
+    });
   });
 
   it('properly generates custom directive values', async () => {

tests/zod.spec.ts

@@ -662,6 +662,7 @@ describe('zod', () => {
         // Shape Schema
         'export function ShapeSchema() {',
         'z.union([CircleSchema(), SquareSchema()])',
+        '}',
       ];
       for (const wantContain of wantContains) {
         expect(result.content).toContain(wantContain);