diff --git a/example/myzod/schemas.ts b/example/myzod/schemas.ts
index 8cac6fdb..36b9d2b7 100644
--- a/example/myzod/schemas.ts
+++ b/example/myzod/schemas.ts
@@ -1,8 +1,15 @@
 import * as myzod from 'myzod'
-import { AttributeInput, ButtonComponentType, ComponentInput, DropDownComponentInput, EventArgumentInput, EventInput, EventOptionType, HttpInput, HttpMethod, LayoutInput, PageInput, PageType, User } from '../types'
+import { Admin, AttributeInput, ButtonComponentType, ComponentInput, DropDownComponentInput, EventArgumentInput, EventInput, EventOptionType, Guest, HttpInput, HttpMethod, LayoutInput, PageInput, PageType, User } from '../types'
 
 export const definedNonNullAnySchema = myzod.object({});
 
+export function AdminSchema(): myzod.Type<Admin> {
+  return myzod.object({
+    __typename: myzod.literal('Admin').optional(),
+    lastModifiedAt: definedNonNullAnySchema.optional().nullable()
+  })
+}
+
 export function AttributeInputSchema(): myzod.Type<AttributeInput> {
   return myzod.object({
     key: myzod.string().optional().nullable(),
@@ -45,6 +52,13 @@ export function EventInputSchema(): myzod.Type<EventInput> {
 
 export const EventOptionTypeSchema = myzod.enum(EventOptionType);
 
+export function GuestSchema(): myzod.Type<Guest> {
+  return myzod.object({
+    __typename: myzod.literal('Guest').optional(),
+    lastLoggedIn: definedNonNullAnySchema.optional().nullable()
+  })
+}
+
 export function HttpInputSchema(): myzod.Type<HttpInput> {
   return myzod.object({
     method: HttpMethodSchema.optional().nullable(),
@@ -84,8 +98,13 @@ export function UserSchema(): myzod.Type<User> {
     createdAt: definedNonNullAnySchema.optional().nullable(),
     email: myzod.string().optional().nullable(),
     id: myzod.string().optional().nullable(),
+    kind: UserKindSchema().optional().nullable(),
     name: myzod.string().optional().nullable(),
     password: myzod.string().optional().nullable(),
     updatedAt: definedNonNullAnySchema.optional().nullable()
   })
 }
+
+export function UserKindSchema() {
+  return myzod.union([AdminSchema(), GuestSchema()])
+}
diff --git a/example/test.graphql b/example/test.graphql
index 93cf53b8..62450491 100644
--- a/example/test.graphql
+++ b/example/test.graphql
@@ -5,11 +5,22 @@ enum PageType {
   BASIC_AUTH
 }
 
+type Admin {
+  lastModifiedAt: Date
+}
+
+type Guest {
+  lastLoggedIn: Date
+}
+
+union UserKind = Admin | Guest
+
 type User {
   id: ID
   name: String
   email: String
   password: String
+  kind: UserKind
   createdAt: Date
   updatedAt: Date
 }
diff --git a/example/types.ts b/example/types.ts
index b0b5bd10..20898898 100644
--- a/example/types.ts
+++ b/example/types.ts
@@ -14,6 +14,11 @@ export type Scalars = {
   URL: any;
 };
 
+export type Admin = {
+  __typename?: 'Admin';
+  lastModifiedAt?: Maybe<Scalars['Date']>;
+};
+
 export type AttributeInput = {
   key?: InputMaybe<Scalars['String']>;
   val?: InputMaybe<Scalars['String']>;
@@ -52,6 +57,11 @@ export enum EventOptionType {
   Retry = 'RETRY'
 }
 
+export type Guest = {
+  __typename?: 'Guest';
+  lastLoggedIn?: Maybe<Scalars['Date']>;
+};
+
 export type HttpInput = {
   method?: InputMaybe<HttpMethod>;
   url: Scalars['URL'];
@@ -92,7 +102,10 @@ export type User = {
   createdAt?: Maybe<Scalars['Date']>;
   email?: Maybe<Scalars['String']>;
   id?: Maybe<Scalars['ID']>;
+  kind?: Maybe<UserKind>;
   name?: Maybe<Scalars['String']>;
   password?: Maybe<Scalars['String']>;
   updatedAt?: Maybe<Scalars['Date']>;
 };
+
+export type UserKind = Admin | Guest;
diff --git a/example/yup/schemas.ts b/example/yup/schemas.ts
index 59d9dbd6..f574b5e3 100644
--- a/example/yup/schemas.ts
+++ b/example/yup/schemas.ts
@@ -1,5 +1,18 @@
 import * as yup from 'yup'
-import { AttributeInput, ButtonComponentType, ComponentInput, DropDownComponentInput, EventArgumentInput, EventInput, EventOptionType, HttpInput, HttpMethod, LayoutInput, PageInput, PageType, User } from '../types'
+import { Admin, AttributeInput, ButtonComponentType, ComponentInput, DropDownComponentInput, EventArgumentInput, EventInput, EventOptionType, Guest, HttpInput, HttpMethod, LayoutInput, PageInput, PageType, User, UserKind } from '../types'
+
+function union<T>(...schemas: ReadonlyArray<yup.SchemaOf<T>>): yup.BaseSchema<T> {
+  return yup.mixed().test({
+    test: (value) => schemas.some((schema) => schema.isValidSync(value))
+  })
+}
+
+export function AdminSchema(): yup.SchemaOf<Admin> {
+  return yup.object({
+    __typename: yup.mixed().oneOf(['Admin', undefined]),
+    lastModifiedAt: yup.mixed()
+  })
+}
 
 export function AttributeInputSchema(): yup.SchemaOf<AttributeInput> {
   return yup.object({
@@ -43,6 +56,13 @@ export function EventInputSchema(): yup.SchemaOf<EventInput> {
 
 export const EventOptionTypeSchema = yup.mixed().oneOf([EventOptionType.Reload, EventOptionType.Retry]);
 
+export function GuestSchema(): yup.SchemaOf<Guest> {
+  return yup.object({
+    __typename: yup.mixed().oneOf(['Guest', undefined]),
+    lastLoggedIn: yup.mixed()
+  })
+}
+
 export function HttpInputSchema(): yup.SchemaOf<HttpInput> {
   return yup.object({
     method: HttpMethodSchema,
@@ -82,8 +102,13 @@ export function UserSchema(): yup.SchemaOf<User> {
     createdAt: yup.mixed(),
     email: yup.string(),
     id: yup.string(),
+    kind: UserKindSchema(),
     name: yup.string(),
     password: yup.string(),
     updatedAt: yup.mixed()
   })
 }
+
+export function UserKindSchema(): yup.BaseSchema<UserKind> {
+  return union<UserKind>(AdminSchema(), GuestSchema())
+}
diff --git a/example/zod/schemas.ts b/example/zod/schemas.ts
index 2d576abe..3b294d28 100644
--- a/example/zod/schemas.ts
+++ b/example/zod/schemas.ts
@@ -1,5 +1,5 @@
 import { z } from 'zod'
-import { AttributeInput, ButtonComponentType, ComponentInput, DropDownComponentInput, EventArgumentInput, EventInput, EventOptionType, HttpInput, HttpMethod, LayoutInput, PageInput, PageType, User } from '../types'
+import { Admin, AttributeInput, ButtonComponentType, ComponentInput, DropDownComponentInput, EventArgumentInput, EventInput, EventOptionType, Guest, HttpInput, HttpMethod, LayoutInput, PageInput, PageType, User } from '../types'
 
 type Properties<T> = Required<{
   [K in keyof T]: z.ZodType<T[K], any, T[K]>;
@@ -11,6 +11,13 @@ export const isDefinedNonNullAny = (v: any): v is definedNonNullAny => v !== und
 
 export const definedNonNullAnySchema = z.any().refine((v) => isDefinedNonNullAny(v));
 
+export function AdminSchema(): z.ZodObject<Properties<Admin>> {
+  return z.object<Properties<Admin>>({
+    __typename: z.literal('Admin').optional(),
+    lastModifiedAt: definedNonNullAnySchema.nullish()
+  })
+}
+
 export function AttributeInputSchema(): z.ZodObject<Properties<AttributeInput>> {
   return z.object<Properties<AttributeInput>>({
     key: z.string().nullish(),
@@ -53,6 +60,13 @@ export function EventInputSchema(): z.ZodObject<Properties<EventInput>> {
 
 export const EventOptionTypeSchema = z.nativeEnum(EventOptionType);
 
+export function GuestSchema(): z.ZodObject<Properties<Guest>> {
+  return z.object<Properties<Guest>>({
+    __typename: z.literal('Guest').optional(),
+    lastLoggedIn: definedNonNullAnySchema.nullish()
+  })
+}
+
 export function HttpInputSchema(): z.ZodObject<Properties<HttpInput>> {
   return z.object<Properties<HttpInput>>({
     method: HttpMethodSchema.nullish(),
@@ -92,8 +106,13 @@ export function UserSchema(): z.ZodObject<Properties<User>> {
     createdAt: definedNonNullAnySchema.nullish(),
     email: z.string().nullish(),
     id: z.string().nullish(),
+    kind: UserKindSchema().nullish(),
     name: z.string().nullish(),
     password: z.string().nullish(),
     updatedAt: definedNonNullAnySchema.nullish()
   })
 }
+
+export function UserKindSchema() {
+  return z.union([AdminSchema(), GuestSchema()])
+}
diff --git a/package.json b/package.json
index d45646e9..e1dae968 100644
--- a/package.json
+++ b/package.json
@@ -16,6 +16,9 @@
   ],
   "scripts": {
     "type-check": "tsc --noEmit",
+    "type-check:yup": "tsc --strict --noEmit example/yup/schemas.ts",
+    "type-check:zod": "tsc --strict --noEmit example/zod/schemas.ts",
+    "type-check:myzod": "tsc --strict --noEmit example/myzod/schemas.ts",
     "test": "jest --no-watchman",
     "build": "run-p build:*",
     "build:main": "tsc -p tsconfig.main.json",
diff --git a/src/myzod/index.ts b/src/myzod/index.ts
index e36abcd6..cee23817 100644
--- a/src/myzod/index.ts
+++ b/src/myzod/index.ts
@@ -91,6 +91,8 @@ export const MyZodSchemaVisitor = (schema: GraphQLSchema, config: ValidationSche
         .withContent(`myzod.enum(${enumname})`).string;
     },
     UnionTypeDefinition: (node: UnionTypeDefinitionNode) => {
+      if (!node.types) return;
+
       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;
@@ -98,13 +100,8 @@ export const MyZodSchemaVisitor = (schema: GraphQLSchema, config: ValidationSche
       const union =
         unionElementsCount > 1 ? indent(`return myzod.union([${unionElements}])`) : indent(`return ${unionElements}`);
 
-      const result = new DeclarationBlock({})
-        .export()
-        .asKind('function')
-        .withName(`${unionName}Schema()`)
-        .withBlock(union);
-
-      return result.string;
+      return new DeclarationBlock({}).export().asKind('function').withName(`${unionName}Schema()`).withBlock(union)
+        .string;
     },
   };
 };
diff --git a/src/yup/index.ts b/src/yup/index.ts
index 401bcf3b..df6da2f1 100644
--- a/src/yup/index.ts
+++ b/src/yup/index.ts
@@ -30,6 +30,7 @@ export const YupSchemaVisitor = (schema: GraphQLSchema, config: ValidationSchema
       return [importYup];
     },
     initialEmit: (): string =>
+      '\n' +
       new DeclarationBlock({})
         .asKind('function')
         .withName('union<T>(...schemas: ReadonlyArray<yup.SchemaOf<T>>): yup.BaseSchema<T>')
@@ -101,22 +102,19 @@ export const YupSchemaVisitor = (schema: GraphQLSchema, config: ValidationSchema
         .withContent(`yup.mixed().oneOf([${values}])`).string;
     },
     UnionTypeDefinition: (node: UnionTypeDefinitionNode) => {
+      if (!node.types) return;
+
       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;
+      importTypes.push(unionName);
 
-      const union =
-        unionElementsCount > 1
-          ? indent(`return union<${unionName}>(${unionElements})`)
-          : indent(`return ${unionElements}`);
+      const unionElements = node.types?.map(t => `${tsVisitor.convertName(t.name.value)}Schema()`).join(', ');
+      const union = indent(`return union<${unionName}>(${unionElements})`);
 
-      const result = new DeclarationBlock({})
+      return new DeclarationBlock({})
         .export()
         .asKind('function')
         .withName(`${unionName}Schema(): yup.BaseSchema<${unionName}>`)
-        .withBlock(union);
-
-      return result.string;
+        .withBlock(union).string;
     },
     // ScalarTypeDefinition: (node) => {
     //   const decl = new DeclarationBlock({})
@@ -209,6 +207,11 @@ const generateNameNodeYupSchema = (
     return `${enumName}Schema`;
   }
 
+  if (typ?.astNode?.kind === 'UnionTypeDefinition') {
+    const enumName = tsVisitor.convertName(typ.astNode.name.value);
+    return `${enumName}Schema()`;
+  }
+
   const primitive = yup4Scalar(config, tsVisitor, node.value);
   return primitive;
 };
diff --git a/src/zod/index.ts b/src/zod/index.ts
index 894b6e6a..ce36ecb8 100644
--- a/src/zod/index.ts
+++ b/src/zod/index.ts
@@ -102,20 +102,17 @@ export const ZodSchemaVisitor = (schema: GraphQLSchema, config: ValidationSchema
         .withContent(`z.nativeEnum(${enumname})`).string;
     },
     UnionTypeDefinition: (node: UnionTypeDefinitionNode) => {
+      if (!node.types) return;
+
       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 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(union);
-
-      return result.string;
+      return new DeclarationBlock({}).export().asKind('function').withName(`${unionName}Schema()`).withBlock(union)
+        .string;
     },
   };
 };
diff --git a/tests/myzod.spec.ts b/tests/myzod.spec.ts
index c4e950cd..f03c2b3c 100644
--- a/tests/myzod.spec.ts
+++ b/tests/myzod.spec.ts
@@ -563,7 +563,7 @@ describe('myzod', () => {
       const wantContains = [
         // Shape Schema
         'export function ShapeSchema() {',
-        'myzod.union([CircleSchema(), SquareSchema()])',
+        'return myzod.union([CircleSchema(), SquareSchema()])',
         '}',
       ];
       for (const wantContain of wantContains) {
@@ -629,7 +629,7 @@ describe('myzod', () => {
       const wantContains = [
         // Shape Schema
         'export function ShapeSchema() {',
-        'CircleSchema()',
+        'return CircleSchema()',
         '}',
       ];
       for (const wantContain of wantContains) {
diff --git a/tests/yup.spec.ts b/tests/yup.spec.ts
index e7697237..a5a81779 100644
--- a/tests/yup.spec.ts
+++ b/tests/yup.spec.ts
@@ -513,7 +513,7 @@ describe('yup', () => {
         'export function GeometrySchema(): yup.SchemaOf<Geometry> {',
         'return yup.object({',
         "__typename: yup.mixed().oneOf(['Geometry', undefined]),",
-        'shape: yup.mixed()',
+        'shape: ShapeSchema()',
         '})',
       ];
       for (const wantContain of wantContains) {
@@ -542,7 +542,7 @@ describe('yup', () => {
       const wantContains = [
         // Shape Schema
         'export function ShapeSchema(): yup.BaseSchema<Shape> {',
-        'CircleSchema()',
+        'return union<Shape>(CircleSchema())',
         '}',
       ];
       for (const wantContain of wantContains) {
diff --git a/tests/zod.spec.ts b/tests/zod.spec.ts
index e6f4556d..c278ce00 100644
--- a/tests/zod.spec.ts
+++ b/tests/zod.spec.ts
@@ -661,7 +661,7 @@ describe('zod', () => {
       const wantContains = [
         // Shape Schema
         'export function ShapeSchema() {',
-        'z.union([CircleSchema(), SquareSchema()])',
+        'return z.union([CircleSchema(), SquareSchema()])',
         '}',
       ];
       for (const wantContain of wantContains) {
@@ -725,7 +725,7 @@ describe('zod', () => {
       const wantContains = [
         // Shape Schema
         'export function ShapeSchema() {',
-        'CircleSchema()',
+        'return CircleSchema()',
         '}',
       ];
       for (const wantContain of wantContains) {