Clean up CLI arg parsing & passing

libs/checkpoint-validation/bin/jest.config.cjs → libs/checkpoint-validation/bin/jest.config.js

@@ -1,20 +1,22 @@
 // This is the Jest config used by the test harness when being executed via the CLI.
 // For the Jest config for the tests in this project, see the `jest.config.cjs` in the root of the package workspace.
-const path = require("path");
+import path from "node:path";
+import { fileURLToPath } from "node:url";
+import { parseArgs } from "../dist/parse_args.js";
+
+const args = await parseArgs(process.argv.slice(2));
 
 /** @type {import('ts-jest').JestConfigWithTsJest} */
-const config = {
+export default {
   preset: "ts-jest/presets/default-esm",
-  rootDir: path.resolve(__dirname, "..", "dist"),
+  rootDir: path.resolve(path.dirname(fileURLToPath(import.meta.url)), "..", "dist"),
   testEnvironment: "node",
   testMatch: ["<rootDir>/runner.js"],
   transform: {
-    "^.+\\.(ts|js)x?$": ["@swc/jest"],
+    "^.+\\.[jt]sx?$": ["@swc/jest"],
   },
   moduleNameMapper: {
     "^(\\.{1,2}/.*)\\.[jt]sx?$": "$1",
   },
-  maxWorkers: "50%",
+  globals: args,
 };
-
-module.exports = config;

libs/checkpoint-validation/src/cli.ts

@@ -1,109 +1,23 @@
 import { dirname, resolve as pathResolve } from "node:path";
 import { fileURLToPath } from "node:url";
 import { runCLI } from "@jest/core";
-import yargs, { ArgumentsCamelCase } from "yargs";
-import { BaseCheckpointSaver } from "@langchain/langgraph-checkpoint";
-import {
-  CheckpointerTestInitializer,
-  checkpointerTestInitializerSchema,
-  GlobalThis,
-  TestTypeFilter,
-} from "./types.js";
-import { dynamicImport, resolveImportPath } from "./import_utils.js";
 
 // make it so we can import/require .ts files
 import "@swc-node/register/esm-register";
+import { parseArgs } from "./parse_args.js";
 
 export async function main() {
   const moduleDirname = dirname(fileURLToPath(import.meta.url));
 
-  const builder = yargs();
-  await builder
-    .command(
-      "* <initializer-import-path> [filters..]",
-      "Validate a checkpointer",
-      {
-        builder: (args) => {
-          return args
-            .positional("initializerImportPath", {
-              type: "string",
-              describe:
-                "The import path of the CheckpointSaverTestInitializer for the checkpointer (passed to 'import()'). " +
-                "Must be the default export.",
-              demandOption: true,
-            })
-            .positional("filters", {
-              array: true,
-              choices: ["getTuple", "put", "putWrites", "list"],
-              default: [],
-              describe:
-                "Only run the specified suites. Valid values are 'getTuple', 'put', 'putWrites', and 'list'",
-              demandOption: false,
-            });
-        },
-        handler: async (
-          argv: ArgumentsCamelCase<{
-            initializerImportPath: string;
-            filters: string[];
-          }>
-        ) => {
-          const { initializerImportPath, filters } = argv;
-
-          let resolvedImportPath;
-
-          try {
-            resolvedImportPath = resolveImportPath(initializerImportPath);
-          } catch (e) {
-            console.error(
-              `Failed to resolve import path '${initializerImportPath}': ${e}`
-            );
-            process.exit(1);
-          }
-
-          let initializerExport: unknown;
-          try {
-            initializerExport = await dynamicImport(resolvedImportPath);
-          } catch (e) {
-            console.error(
-              `Failed to import initializer from import path '${initializerImportPath}' (resolved to '${resolvedImportPath}'): ${e}`
-            );
-            process.exit(1);
-          }
-
-          let initializer: CheckpointerTestInitializer<BaseCheckpointSaver>;
-          try {
-            initializer = checkpointerTestInitializerSchema.parse(
-              (initializerExport as { default?: unknown }).default ??
-                initializerExport
-            );
-            (
-              globalThis as GlobalThis
-            ).__langgraph_checkpoint_validation_initializer = initializer;
-            (
-              globalThis as GlobalThis
-            ).__langgraph_checkpoint_validation_filters =
-              filters as TestTypeFilter[];
-          } catch (e) {
-            console.error(
-              `Initializer imported from '${initializerImportPath}' does not conform to the expected schema. Make sure " +
-              "it is the default export, and that implements the CheckpointSaverTestInitializer interface. Error: ${e}`
-            );
-            process.exit(1);
-          }
-
-          await runCLI(
-            {
-              _: [],
-              $0: "",
-            },
-            [pathResolve(moduleDirname, "..", "bin", "jest.config.cjs")]
-          );
-        },
-      }
-    )
-    .help()
-    .alias("h", "help")
-    .wrap(builder.terminalWidth())
-    .strict()
-    .parseAsync(process.argv.slice(2));
+  // parse here to check for errors before running Jest
+  await parseArgs(process.argv.slice(2));
+
+  await runCLI(
+    {
+      _: [],
+      $0: "",
+      runInBand: true,
+    },
+    [pathResolve(moduleDirname, "..", "bin", "jest.config.js")]
+  );
 }

libs/checkpoint-validation/src/parse_args.ts

@@ -0,0 +1,103 @@
+import yargs from "yargs";
+import { BaseCheckpointSaver } from "@langchain/langgraph-checkpoint";
+import {
+  CheckpointerTestInitializer,
+  checkpointerTestInitializerSchema,
+  isTestTypeFilter,
+  isTestTypeFilterArray,
+  TestTypeFilter,
+  testTypeFilters,
+} from "./types.js";
+import { dynamicImport, resolveImportPath } from "./import_utils.js";
+
+// We have to Symbol.for here instead of unique symbols because jest gives each test file a unique module cache, so
+// these symbols get created once when the jest config is evaluated, and again when runner.ts executes, making it
+// impossible for runner.ts to use them as global keys.
+const symbolPrefix = "langgraph-checkpoint-validation";
+export const initializerSymbol = Symbol.for(`${symbolPrefix}-initializer`);
+export const filtersSymbol = Symbol.for(`${symbolPrefix}-filters`);
+
+export type ParsedArgs<
+  CheckpointerT extends BaseCheckpointSaver = BaseCheckpointSaver
+> = {
+  [initializerSymbol]: CheckpointerTestInitializer<CheckpointerT>;
+  [filtersSymbol]: TestTypeFilter[];
+};
+
+const builder = yargs()
+  .command("* <initializer-import-path> [filters..]", "Validate a checkpointer")
+  .positional("initializerImportPath", {
+    type: "string",
+    describe:
+      "The import path of the CheckpointSaverTestInitializer for the checkpointer (passed to 'import()'). " +
+      "Must be the default export.",
+    demandOption: true,
+  })
+  .positional("filters", {
+    array: true,
+    choices: ["getTuple", "put", "putWrites", "list"],
+    default: [],
+    describe:
+      "Only run the specified suites. Valid values are 'getTuple', 'put', 'putWrites', and 'list'",
+    demandOption: false,
+  })
+  .help()
+  .alias("h", "help")
+  .wrap(yargs().terminalWidth())
+  .strict();
+
+export async function parseArgs<CheckpointerT extends BaseCheckpointSaver>(
+  argv: string[]
+): Promise<ParsedArgs<CheckpointerT>> {
+  const { initializerImportPath, filters } = await builder.parse(argv);
+
+  let resolvedImportPath;
+
+  try {
+    resolvedImportPath = resolveImportPath(initializerImportPath);
+  } catch (e) {
+    console.error(
+      `Failed to resolve import path '${initializerImportPath}': ${e}`
+    );
+    process.exit(1);
+  }
+
+  let initializerExport: unknown;
+  try {
+    initializerExport = await dynamicImport(resolvedImportPath);
+  } catch (e) {
+    console.error(
+      `Failed to import initializer from import path '${initializerImportPath}' (resolved to '${resolvedImportPath}'): ${e}`
+    );
+    process.exit(1);
+  }
+
+  let initializer: CheckpointerTestInitializer<CheckpointerT>;
+  try {
+    initializer = checkpointerTestInitializerSchema.parse(
+      (initializerExport as { default?: unknown }).default ?? initializerExport
+    ) as CheckpointerTestInitializer<CheckpointerT>;
+  } catch (e) {
+    console.error(
+      `Initializer imported from '${initializerImportPath}' does not conform to the expected schema. Make sure " +
+      "it is the default export, and that implements the CheckpointSaverTestInitializer interface. Error: ${e}`
+    );
+    process.exit(1);
+  }
+
+  if (!isTestTypeFilterArray(filters)) {
+    console.error(
+      `Invalid filters: '${filters
+        .filter((f) => !isTestTypeFilter(f))
+        .join("', '")}'. Expected only values from '${testTypeFilters.join(
+        "', '"
+      )}'`
+    );
+    process.exit(1);
+  }
+
+  return {
+    [initializerSymbol]: initializer,
+    [filtersSymbol]: filters,
+  };
+}

libs/checkpoint-validation/src/runner.ts

@@ -2,19 +2,17 @@
 // Jest test file because unfortunately there's no good way to just pass Jest a test definition function and tell it to
 // run it.
 import { specTest } from "./spec/index.js";
-import type { GlobalThis } from "./types.js";
+import { ParsedArgs, filtersSymbol, initializerSymbol } from "./parse_args.js";
 
 // passing via global is ugly, but there's no good alternative for handling the dynamic import here
-const initializer = (globalThis as GlobalThis)
-  .__langgraph_checkpoint_validation_initializer;
+const initializer = (globalThis as typeof globalThis & ParsedArgs)[
+  initializerSymbol
+];
 
 if (!initializer) {
-  throw new Error(
-    "expected global '__langgraph_checkpoint_validation_initializer' is not set"
-  );
+  throw new Error("Test configuration error: initializer is not set.");
 }
 
-const filters = (globalThis as GlobalThis)
-  .__langgraph_checkpoint_validation_filters;
+const filters = (globalThis as typeof globalThis & ParsedArgs)[filtersSymbol];
 
 specTest(initializer, filters);

libs/checkpoint-validation/src/types.ts

@@ -70,9 +70,21 @@ export const checkpointerTestInitializerSchema = z.object({
     .optional(),
 });
 
-export type TestTypeFilter = "getTuple" | "list" | "put" | "putWrites";
+export const testTypeFilters = [
+  "getTuple",
+  "list",
+  "put",
+  "putWrites",
+] as const;
 
-export type GlobalThis = typeof globalThis & {
-  __langgraph_checkpoint_validation_initializer?: CheckpointerTestInitializer<BaseCheckpointSaver>;
-  __langgraph_checkpoint_validation_filters?: TestTypeFilter[];
-};
+export type TestTypeFilter = (typeof testTypeFilters)[number];
+
+export function isTestTypeFilter(value: string): value is TestTypeFilter {
+  return testTypeFilters.includes(value as TestTypeFilter);
+}
+
+export function isTestTypeFilterArray(
+  value: string[]
+): value is TestTypeFilter[] {
+  return value.every(isTestTypeFilter);
+}

libs/checkpoint-validation/tsconfig.cjs.json

@@ -11,7 +11,8 @@
     "docs",
     "**/tests",
     "src/cli.ts",
-    "src/importUtils.ts",
-    "src/runner.ts"
+    "src/import_utils.ts",
+    "src/runner.ts",
+    "src/parse_args.ts"
   ]
 }