Merge pull request #55 from josh-degraw/discriminated-union

Add support for discriminated unions
anatine · Sep 5, 2022 · 399afe3 · 399afe3
2 parents 40d2d1f + 0dbcf9e
commit 399afe3
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 6 deletions.
diff --git a/packages/zod-mock/src/lib/zod-mock.spec.ts b/packages/zod-mock/src/lib/zod-mock.spec.ts
@@ -21,7 +21,11 @@ describe('zod-mock', () => {
       numberCount: z.number().transform((item) => `total value = ${item}`),
       age: z.number().min(18).max(120),
       record: z.record(z.string(), z.number()),
-      nativeEnum: z.nativeEnum(NativeEnum)
+      nativeEnum: z.nativeEnum(NativeEnum),
+      discriminatedUnion: z.discriminatedUnion('discriminator', [
+        z.object({ discriminator: z.literal('a'), a: z.boolean() }),
+        z.object({ discriminator: z.literal('b'), b: z.string() }),
+      ]),
     });
 
     const mockData = generateMock(schema); //?
@@ -40,10 +44,12 @@ describe('zod-mock', () => {
     expect(typeof mockData.jobTitle).toEqual('string');
     expect(typeof mockData.stringLength).toEqual('number');
     expect(typeof mockData.numberCount).toEqual('string');
-    expect(mockData.age >= 18 && mockData.age <= 120).toBeTruthy();
+    expect(mockData.age).toBeGreaterThanOrEqual(18);
+    expect(mockData.age).toBeLessThanOrEqual(120);
     expect(typeof mockData.record).toEqual('object');
     expect(typeof Object.values(mockData.record)[0]).toEqual('number');
     expect(mockData.nativeEnum === 1 || mockData.nativeEnum === 2);
+    expect(mockData.discriminatedUnion).toBeTruthy()
   });
 
   it('should generate mock data of the appropriate type when the field names overlap Faker properties that are not valid functions', () => {
@@ -359,4 +365,27 @@ describe('zod-mock', () => {
     expect(transformResult.items).toBeGreaterThan(0);
     expect(transformResult.items).toBeLessThan(101);
   });
+
+  it('should handle discriminated unions', () => {
+    const FirstType = z.object({
+      hasEmail: z.literal(false),
+      userName: z.string(),
+    });
+
+    const SecondType = z.object({
+      hasEmail: z.literal(true),
+      email: z.string(),
+    });
+
+    const Union = z.discriminatedUnion('hasEmail', [FirstType, SecondType]);
+
+    const result = generateMock(Union);
+    expect(result).toBeDefined();
+
+    if (result.hasEmail) {
+      expect(result.email).toBeTruthy();
+    } else {
+      expect(result.userName).toBeTruthy();
+    }
+  });
 });
diff --git a/packages/zod-mock/src/lib/zod-mock.ts b/packages/zod-mock/src/lib/zod-mock.ts
@@ -290,11 +290,25 @@ function parseEnum(zodRef: z.ZodEnum<never> | z.ZodNativeEnum<never>) {
   return values[pick];
 }
 
+function parseDiscriminatedUnion(
+  zodRef: z.ZodDiscriminatedUnion<never, never, never>,
+  options?: GenerateMockOptions
+) {
+
+  // Map the options to various possible union cases
+  const potentialCases = [...zodRef._def.options.values()];
+  const pick = Math.floor(Math.random() * potentialCases.length);
+
+  const mocked = potentialCases[pick];
+
+  return generateMock(mocked, options);
+}
+
 function parseNativeEnum(zodRef: z.ZodNativeEnum<never>) {
-	const { values } = zodRef._def;
-	const pick = Math.floor(Math.random() * (Object.values(values).length));
-	const key = Array.from(Object.keys(values))[pick];
-	return values[values[key]];
+  const { values } = zodRef._def;
+  const pick = Math.floor(Math.random() * Object.values(values).length);
+  const key = Array.from(Object.keys(values))[pick];
+  return values[values[key]];
 }
 
 function parseLiteral(zodRef: z.ZodLiteral<any>) {
@@ -343,6 +357,7 @@ const workerMap = {
   ZodTransformer: parseTransform,
   ZodEffects: parseTransform,
   ZodUnion: parseUnion,
+  ZodDiscriminatedUnion: parseDiscriminatedUnion,
 };
 type WorkerKeys = keyof typeof workerMap;