[zod-mock]: adds support for specifying ZodDiscriminatedUnion variant when discriminator value is set in stringMap options

packages/zod-mock/README.md

@@ -90,6 +90,8 @@ const mockData = generateMock(schema, {
 });
 ```
 
+If a supplied string value maps to a `Discriminator` field for a `ZodDiscriminatedUnion`, the generator will produce a mock of the appropriate discriminated variant.
+
 ----
 
 ## Adding a seed generator

packages/zod-mock/src/lib/zod-mock.spec.ts

@@ -594,27 +594,91 @@ describe('zod-mock', () => {
     expect(transformResult.items).toBeLessThan(101);
   });
 
-  it('should handle discriminated unions', () => {
-    const FirstType = z.object({
-      hasEmail: z.literal(false),
-      userName: z.string(),
-    });
+  describe('ZodDiscriminatedUnion', () => {
+
+    it('should handle discriminated unions', () => {
 
-    const SecondType = z.object({
-      hasEmail: z.literal(true),
-      email: z.string(),
+      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();
+      }
     });
 
-    const Union = z.discriminatedUnion('hasEmail', [FirstType, SecondType]);
 
-    const result = generateMock(Union);
-    expect(result).toBeDefined();
+    describe('when the discriminator value is specified in stringMap options', () => {
+      const FirstType = z.object({
+        userType: z.literal('type-1'),
+        userName: z.string(),
+      });
 
-    if (result.hasEmail) {
-      expect(result.email).toBeTruthy();
-    } else {
-      expect(result.userName).toBeTruthy();
-    }
+      const SecondType = z.object({
+        userType: z.literal('type-2'),
+        email: z.string(),
+      });
+
+      const Union = z.discriminatedUnion('userType', [FirstType, SecondType]);
+
+      it('should generate the variant defined by the discriminator value', () => {
+        // validate that we always generate 'type-1' when specified
+        [...Array(10).keys()].forEach(() => {
+          const result = generateMock(Union, { stringMap: {
+            userType: () => 'type-1',
+          }});
+          expect(result).toBeDefined();
+          expect(result).toEqual(expect.objectContaining({
+            userType: 'type-1',
+            userName: expect.any(String),
+          }));
+          expect(result).not.toEqual(expect.objectContaining({
+            email: expect.anything(),
+          }));
+        });
+
+        // validate that we always generate 'type-2' when specified
+        [...Array(10).keys()].forEach(() => {
+          const result = generateMock(Union, { stringMap: {
+            userType: () => 'type-2',
+          }});
+          expect(result).toBeDefined();
+          expect(result).toEqual(expect.objectContaining({
+            userType: 'type-2',
+            email: expect.any(String),
+          }));
+          expect(result).not.toEqual(expect.objectContaining({
+            userName: expect.anything(),
+          }));
+        });
+      });
+
+      it('should fall back to random variant generation if the desired discriminator value is not present in the schema', () => {
+        let result;
+        expect(() => {
+          result = generateMock(Union, { stringMap: {
+            userType: () => 'type-3',
+          }});
+        }).not.toThrow();
+
+        expect(result).toBeDefined();
+        // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+        expect(['type-1', 'type-2']).toContain(result!.userType);
+      });
+    });
   });
 
   it('should handle branded types', () => {

packages/zod-mock/src/lib/zod-mock.ts

@@ -8,6 +8,7 @@ import {
   ZodType,
   ZodString,
   ZodRecord,
+  ZodDiscriminatedUnionOption,
   util,
 } from 'zod';
 import {
@@ -426,6 +427,23 @@ function parseDiscriminatedUnion(
   const fakerInstance = options?.faker || faker;
   // Map the options to various possible union cases
   const potentialCases = [...zodRef._def.options.values()];
+  const discriminatorField = zodRef._def.discriminator as string;
+  const hasSpecifiedDiscrimination = typeof options?.stringMap === 'object' && Object.prototype.hasOwnProperty.call(options?.stringMap, discriminatorField);
+
+  // If a discriminator field value is specified in the stringMap options, generate mock values in alignment with that discriminated shape
+  if (hasSpecifiedDiscrimination) {
+    const desiredDescriminationValueFunction = options?.stringMap?.[discriminatorField];
+    const desiredValue = typeof desiredDescriminationValueFunction === 'function' ? desiredDescriminationValueFunction() : undefined;
+    if (desiredValue !== undefined) {
+      const mocked = (zodRef._def.options as ZodDiscriminatedUnionOption<typeof desiredValue>[]).find((option) => {
+          return option._def.shape()[discriminatorField]?._def?.value === desiredValue;
+      })
+      // if the desired discriminated variant is not found, fall back to random generation
+      if (mocked !== undefined) {
+        return generateMock(mocked as ZodTypeAny, options);
+      }
+    }
+  }
   const mocked = fakerInstance.helpers.arrayElement(potentialCases);
   return generateMock(mocked, options);
 }