refactor: Simplify unique key handling in oneOf function and update related types. Remove itemId from uniqueBy interface, streamline tests to reflect changes in option structure, and ensure compatibility with unique selection logic.

Author: Michal Warda

packages/torque/src/schema-rng.ts

@@ -4,30 +4,29 @@ import { getCurrentMessageContext } from "./message-context";
 import { hasValueBeenUsed, markValueAsUsed } from "./unique-selection";
 import { random } from "./utils";
 
-type PrimitiveKey = string | number | boolean;
-type UniqueIdentifierResolver<T> =
-  | keyof any
-  | ((value: T) => PrimitiveKey | null | undefined);
-
-export interface OneOfUniqueBy<T> {
+export interface OneOfUniqueBy {
   collection: string;
-  itemId?: UniqueIdentifierResolver<T>;
 }
 
 export interface OneOfOptions<T> {
-  unique?: boolean;
-  uniqueBy?: OneOfUniqueBy<T>;
+  uniqueBy?: OneOfUniqueBy;
 }
 
 type NormalizedWeightedOption<T> = {
   value: T;
   weight?: number;
+  id?: string | number | boolean;
   uniqueKey?: string;
 };
 
-type ResolvedUniqueBy<T> = {
+type ResolvedUniqueBy = {
   collection: string;
-  itemId: UniqueIdentifierResolver<T>;
+};
+
+export type OneOfOptionWithId<T> = {
+  value: T;
+  id: string | number | boolean;
+  weight?: number;
 };
 
 export function optional<T>(message: T): T | (() => []) {
@@ -59,6 +58,15 @@ export function randomSample<T>(n: number, array: T[]): T[] {
 
   return result;
 }
+
+export function oneOf<T>(
+  options: Array<OneOfOptionWithId<T>>,
+  config: { uniqueBy: OneOfUniqueBy }
+): T;
+export function oneOf<T>(
+  options: Array<WeightedOneOfOption<T>>,
+  config?: { uniqueBy?: never }
+): T;
 export function oneOf<T>(
   options: Array<WeightedOneOfOption<T>>,
   config?: OneOfOptions<T>
@@ -67,16 +75,8 @@ export function oneOf<T>(
     throw new Error("oneOf requires at least one option");
   }
 
-  const rawUniqueBy = config?.uniqueBy;
-  const enforceUnique = Boolean(config?.unique ?? rawUniqueBy);
-
-  if (enforceUnique && !rawUniqueBy) {
-    throw new Error(
-      "oneOf unique mode requires a uniqueBy option with a collection name"
-    );
-  }
-
-  const uniqueBy = resolveUniqueBy(rawUniqueBy);
+  const uniqueBy = resolveUniqueBy(config?.uniqueBy);
+  const enforceUnique = Boolean(uniqueBy);
 
   const normalized = options.map((option) =>
     isWeightedOption(option) ? { ...option } : { value: option }
@@ -123,18 +123,21 @@ export function oneOf<T>(
     }
   }
 
-  const candidateBase = normalized.map<NormalizedWeightedOption<T>>((option) => ({
-    value: option.value,
-    weight: option.weight,
-  }));
+  const candidateBase = normalized.map<NormalizedWeightedOption<T>>(
+    (option) => ({
+      value: option.value,
+      weight: option.weight,
+      id: option.id,
+    })
+  );
 
   let candidateOptions = candidateBase;
 
   if (enforceUnique && uniqueBy) {
     candidateOptions = candidateBase
       .map((option) => ({
         ...option,
-        uniqueKey: buildUniqueKey(option.value, uniqueBy),
+        uniqueKey: buildUniqueKey(option.id),
       }))
       .filter(
         (option) => !hasValueBeenUsed(uniqueBy.collection, option.uniqueKey!)
@@ -185,7 +188,7 @@ export function oneOf<T>(
 
 function recordUniqueSelection<T>(
   option: NormalizedWeightedOption<T>,
-  uniqueBy: ResolvedUniqueBy<T> | undefined,
+  uniqueBy: ResolvedUniqueBy | undefined,
   enforceUnique: boolean
 ): void {
   if (!enforceUnique || !uniqueBy || !option.uniqueKey) {
@@ -200,51 +203,26 @@ function recordUniqueSelection<T>(
   markValueAsUsed(uniqueBy.collection, option.uniqueKey);
 }
 
-function buildUniqueKey<T>(
-  value: T,
-  uniqueBy: ResolvedUniqueBy<T>
-): string {
-  const rawValue =
-    typeof uniqueBy.itemId === "function"
-      ? uniqueBy.itemId(value)
-      : readProperty(value, uniqueBy.itemId);
-
-  if (
-    rawValue === undefined ||
-    rawValue === null ||
-    (typeof rawValue !== "string" &&
-      typeof rawValue !== "number" &&
-      typeof rawValue !== "boolean")
-  ) {
-    throw new Error(
-      `oneOf uniqueBy.itemId "${String(
-        uniqueBy.itemId
-      )}" must resolve to a string, number, or boolean`
-    );
-  }
-
-  const prefix = typeof rawValue;
-  return `${prefix}:${String(rawValue)}`;
-}
-
-function readProperty<T>(value: T, key: keyof any): unknown {
+function buildUniqueKey(id: unknown): string {
   if (
-    value === null ||
-    (typeof value !== "object" && typeof value !== "function")
+    id === undefined ||
+    id === null ||
+    (typeof id !== "string" &&
+      typeof id !== "number" &&
+      typeof id !== "boolean")
   ) {
     throw new Error(
-      `oneOf uniqueBy.itemId "${String(
-        key
-      )}" requires the option value to be an object or function`
+      `oneOf uniqueBy requires options to be objects with a valid "id" property (string, number, or boolean)`
     );
   }
 
-  return (value as any)[key];
+  const prefix = typeof id;
+  return `${prefix}:${String(id)}`;
 }
 
-function resolveUniqueBy<T>(
-  uniqueBy?: OneOfUniqueBy<T>
-): ResolvedUniqueBy<T> | undefined {
+function resolveUniqueBy(
+  uniqueBy?: OneOfUniqueBy
+): ResolvedUniqueBy | undefined {
   if (!uniqueBy) {
     return undefined;
   }
@@ -258,10 +236,7 @@ function resolveUniqueBy<T>(
     throw new Error("oneOf uniqueBy.collection must be a non-empty string");
   }
 
-  const itemId = uniqueBy.itemId ?? ("id" as UniqueIdentifierResolver<T>);
-
   return {
     collection,
-    itemId,
   };
 }

packages/torque/src/schema.test.ts

@@ -71,22 +71,22 @@ describe("oneOf", () => {
 
   it("enforces uniqueBy across calls", async () => {
     const options = [
-      { value: { id: "a" } },
-      { value: { id: "b" } },
-      { value: { id: "c" } },
+      { id: "a", value: "a" },
+      { id: "b", value: "b" },
+      { id: "c", value: "c" },
     ];
 
     await runWithUniqueSelectionScope(async () => {
       const picks = [
         oneOf([...options], {
-          uniqueBy: { collection: "tools", itemId: "id" },
-        }).id,
+          uniqueBy: { collection: "tools" },
+        }),
         oneOf([...options], {
-          uniqueBy: { collection: "tools", itemId: "id" },
-        }).id,
+          uniqueBy: { collection: "tools" },
+        }),
         oneOf([...options], {
-          uniqueBy: { collection: "tools", itemId: "id" },
-        }).id,
+          uniqueBy: { collection: "tools" },
+        }),
       ];
 
       expect(new Set(picks).size).toBe(3);
@@ -95,68 +95,49 @@ describe("oneOf", () => {
 
   it("throws when a unique collection is exhausted", async () => {
     const options = [
-      { value: { id: "only" } },
-      { value: { id: "only" } },
+      { id: "only", value: "only" },
+      { id: "only", value: "only" },
     ] as const;
 
     await runWithUniqueSelectionScope(async () => {
       oneOf([...options], {
-        uniqueBy: { collection: "limited", itemId: "id" },
+        uniqueBy: { collection: "limited" },
       });
 
       expect(() =>
         oneOf([...options], {
-          uniqueBy: { collection: "limited", itemId: "id" },
+          uniqueBy: { collection: "limited" },
         })
       ).toThrow('oneOf uniqueBy collection "limited" is exhausted');
     });
   });
 
-  it("supports functional unique key extractors", async () => {
+  it("works with uniqueBy and weight together", async () => {
     const options = [
-      { value: { tool: { name: "weather" } } },
-      { value: { tool: { name: "calendar" } } },
+      { id: "a", value: "A", weight: 0.5 },
+      { id: "b", value: "B", weight: 0.3 },
+      { id: "c", value: "C", weight: 0.2 },
     ];
 
     await runWithUniqueSelectionScope(async () => {
-      const first = oneOf([...options], {
-        uniqueBy: {
-          collection: "tools",
-          itemId: (value) => value.tool.name,
-        },
+      await withSeed(100, async () => {
+        const first = oneOf([...options], {
+          uniqueBy: { collection: "weighted-unique" },
+        });
+        expect(["A", "B", "C"]).toContain(first);
+
+        const second = oneOf([...options], {
+          uniqueBy: { collection: "weighted-unique" },
+        });
+        expect(["A", "B", "C"]).toContain(second);
+        expect(second).not.toBe(first);
+
+        const third = oneOf([...options], {
+          uniqueBy: { collection: "weighted-unique" },
+        });
+        expect(["A", "B", "C"]).toContain(third);
+        expect(new Set([first, second, third]).size).toBe(3);
       });
-
-      expect(first.tool.name).toMatch(/weather|calendar/);
-
-      const second = oneOf([...options], {
-        uniqueBy: {
-          collection: "tools",
-          itemId: (value) => value.tool.name,
-        },
-      });
-
-      expect(second.tool.name).not.toBe(first.tool.name);
-    });
-  });
-
-  it("requires uniqueBy when unique is enabled", () => {
-    expect(() =>
-      oneOf([{ value: "a" }, { value: "b" }], { unique: true })
-    ).toThrow("oneOf unique mode requires a uniqueBy option");
-  });
-
-  it("defaults uniqueBy itemId to 'id' when omitted", async () => {
-    const options = [{ value: { id: "alpha" } }, { value: { id: "beta" } }];
-
-    await runWithUniqueSelectionScope(async () => {
-      const first = oneOf([...options], {
-        uniqueBy: { collection: "tools" },
-      }).id;
-      const second = oneOf([...options], {
-        uniqueBy: { collection: "tools" },
-      }).id;
-
-      expect(new Set([first, second]).size).toBe(2);
     });
   });
 });

packages/torque/src/schema.ts

@@ -430,11 +430,12 @@ export type WeightedOneOfOption<T> =
   | {
       value: T;
       weight?: number;
+      id?: string | number | boolean;
     };
 
 export function isWeightedOption<T>(
   option: WeightedOneOfOption<T>
-): option is { value: T; weight?: number } {
+): option is { value: T; weight?: number; id?: string | number | boolean } {
   return typeof option === "object" && option !== null && "value" in option;
 }