Add support for unions in yup

Author: alisabzevari

src/yup/index.ts

@@ -9,6 +9,7 @@ import {
   EnumTypeDefinitionNode,
   ObjectTypeDefinitionNode,
   FieldDefinitionNode,
+  UnionTypeDefinitionNode,
 } from 'graphql';
 import { DeclarationBlock, indent } from '@graphql-codegen/visitor-plugin-common';
 import { TsVisitor } from '@graphql-codegen/typescript';
@@ -28,7 +29,17 @@ export const YupSchemaVisitor = (schema: GraphQLSchema, config: ValidationSchema
       }
       return [importYup];
     },
-    initialEmit: (): string => '',
+    initialEmit: (): string =>
+      new DeclarationBlock({})
+        .asKind('function')
+        .withName('union<T>(...schemas: ReadonlyArray<yup.SchemaOf<T>>): yup.BaseSchema<T>')
+        .withBlock(
+          [
+            indent('return yup.mixed().test({'),
+            indent('test: (value) => schemas.some((schema) => schema.isValidSync(value))', 2),
+            indent('})'),
+          ].join('\n')
+        ).string,
     InputObjectTypeDefinition: (node: InputObjectTypeDefinitionNode) => {
       const name = tsVisitor.convertName(node.name.value);
       importTypes.push(name);
@@ -89,6 +100,24 @@ export const YupSchemaVisitor = (schema: GraphQLSchema, config: ValidationSchema
         .withName(`${enumname}Schema`)
         .withContent(`yup.mixed().oneOf([${values}])`).string;
     },
+    UnionTypeDefinition: (node: UnionTypeDefinitionNode) => {
+      const unionName = tsVisitor.convertName(node.name.value);
+      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 union<${unionName}>(${unionElements})`)
+          : indent(`return ${unionElements}`);
+
+      const result = new DeclarationBlock({})
+        .export()
+        .asKind('function')
+        .withName(`${unionName}Schema(): yup.BaseSchema<${unionName}>`)
+        .withBlock(union);
+
+      return result.string;
+    },
     // ScalarTypeDefinition: (node) => {
     //   const decl = new DeclarationBlock({})
     //     .export()

tests/yup.spec.ts

@@ -451,6 +451,104 @@ describe('yup', () => {
         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: 'yup',
+          withObjectType: true,
+        },
+        {}
+      );
+
+      const wantContains = [
+        // Shape Schema
+        'export function ShapeSchema(): yup.BaseSchema<Shape> {',
+        'union<Shape>(CircleSchema(), SquareSchema())',
+        '}',
+      ];
+      for (const wantContain of wantContains) {
+        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: 'yup',
+          withObjectType: true,
+        },
+        {}
+      );
+
+      const wantContains = [
+        'export function GeometrySchema(): yup.SchemaOf<Geometry> {',
+        'return yup.object({',
+        "__typename: yup.mixed().oneOf(['Geometry', undefined]),",
+        'shape: yup.mixed()',
+        '})',
+      ];
+      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: 'yup',
+          withObjectType: true,
+        },
+        {}
+      );
+
+      const wantContains = [
+        // Shape Schema
+        'export function ShapeSchema(): yup.BaseSchema<Shape> {',
+        'CircleSchema()',
+        '}',
+      ];
+      for (const wantContain of wantContains) {
+        expect(result.content).toContain(wantContain);
+      }
+    });
   });
 
   it('properly generates custom directive values', async () => {