From d189bd44450da95fbcf75f30e54aeb3c93c7a520 Mon Sep 17 00:00:00 2001
From: John Conley <8932043+jfrconley@users.noreply.github.com>
Date: Tue, 12 Mar 2024 14:09:27 -0700
Subject: [PATCH] added support for registry merging and factories

---
 .changeset/light-gifts-shake.md               |   5 +
 .../__tests__/src/attachment-registry.spec.ts | 106 ++++++++
 packages/core/__tests__/src/nornir.spec.ts    | 240 +++++++++++++++---
 packages/core/src/index.ts                    |   2 +-
 packages/core/src/lib/attachment-registry.ts  |  66 ++++-
 packages/core/src/lib/nornir.ts               |   9 +-
 6 files changed, 390 insertions(+), 38 deletions(-)
 create mode 100644 .changeset/light-gifts-shake.md
 create mode 100644 packages/core/__tests__/src/attachment-registry.spec.ts

diff --git a/.changeset/light-gifts-shake.md b/.changeset/light-gifts-shake.md
new file mode 100644
index 0000000..dfc49f3
--- /dev/null
+++ b/.changeset/light-gifts-shake.md
@@ -0,0 +1,5 @@
+---
+"@nornir/core": minor
+---
+
+Added support for registry factories and merging
diff --git a/packages/core/__tests__/src/attachment-registry.spec.ts b/packages/core/__tests__/src/attachment-registry.spec.ts
new file mode 100644
index 0000000..564677c
--- /dev/null
+++ b/packages/core/__tests__/src/attachment-registry.spec.ts
@@ -0,0 +1,106 @@
+import { describe, jest } from "@jest/globals";
+import { AttachmentKey, AttachmentRegistry, RegistryFactory } from "../../dist/lib/attachment-registry.js";
+import { NornirMissingAttachmentException } from "../../dist/lib/error.js";
+
+describe("AttachmentRegistry", () => {
+  let registry: AttachmentRegistry;
+
+  beforeEach(() => {
+    registry = new AttachmentRegistry();
+  });
+
+  describe("#register", () => {
+    it("registers a value and returns a key", () => {
+      const key = registry.register("value");
+      expect(registry.get(key)).toBe("value");
+    });
+  });
+
+  describe("#registerFactory", () => {
+    it("registers a factory and returns a key", () => {
+      const factory: RegistryFactory<string> = jest.fn(() => "value");
+      const key = registry.registerFactory(factory);
+      expect(registry.get(key)).toBe("value");
+    });
+
+    it("caches the result of the factory", () => {
+      const factory: RegistryFactory<string> = jest.fn(() => "value");
+      const key = registry.registerFactory(factory);
+      expect(registry.get(key)).toBe("value");
+      expect(registry.get(key)).toBe("value");
+      expect(factory).toHaveBeenCalledTimes(1);
+    });
+  });
+
+  describe("#put", () => {
+    it("puts a constant value in the registry", () => {
+      const key = AttachmentRegistry.createKey<string>();
+      registry.put(key, "value");
+      expect(registry.get(key)).toBe("value");
+    });
+  });
+
+  describe("#putFactory", () => {
+    it("puts a factory in the registry", () => {
+      const factory: RegistryFactory<string> = jest.fn(() => "value");
+      const key = AttachmentRegistry.createKey<string>();
+      registry.putFactory(key, factory);
+      expect(registry.get(key)).toBe("value");
+    });
+  });
+
+  describe("#delete", () => {
+    it("deletes value from the registry by key", () => {
+      const key = registry.register("value");
+      registry.delete(key);
+      expect(registry.get(key)).toBeNull();
+    });
+  });
+
+  describe("#getAssert", () => {
+    it("gets a value from the registry or throws", () => {
+      const key = registry.register("value");
+      expect(registry.getAssert(key)).toBe("value");
+    });
+
+    it("throws NornirMissingAttachmentException if key does not exist in registry", () => {
+      const key = AttachmentRegistry.createKey<string>();
+      expect(() => registry.getAssert(key)).toThrow(NornirMissingAttachmentException);
+    });
+  });
+
+  describe("#has", () => {
+    it("checks if a key exists in the registry", () => {
+      const key = registry.register("value");
+      expect(registry.has(key)).toBe(true);
+      registry.delete(key);
+      expect(registry.has(key)).toBe(false);
+    });
+  });
+
+  describe("#hasAll", () => {
+    it("checks if all keys exist in the registry", () => {
+      const keys = [registry.register("value1"), registry.register("value2")];
+      expect(registry.hasAll(keys)).toBe(true);
+      registry.delete(keys[0]);
+      expect(registry.hasAll(keys)).toBe(false);
+    });
+  });
+
+  describe("#hasAny", () => {
+    it("checks if any key exists in the registry", () => {
+      const keys = [registry.register("value1"), AttachmentRegistry.createKey<string>()];
+      expect(registry.hasAny(keys)).toBe(true);
+    });
+  });
+
+  describe("#merge", () => {
+    it("merges two registries", () => {
+      const key = registry.register("value");
+      const other = new AttachmentRegistry();
+      other.put(key, "other");
+      registry.merge(other);
+      expect(registry.get(key)).toBe("other");
+    });
+  });
+});
diff --git a/packages/core/__tests__/src/nornir.spec.ts b/packages/core/__tests__/src/nornir.spec.ts
index 95dcc83..fcc1071 100644
--- a/packages/core/__tests__/src/nornir.spec.ts
+++ b/packages/core/__tests__/src/nornir.spec.ts
@@ -1,29 +1,53 @@
 import { describe, expect, jest } from "@jest/globals";
-import { nornir, Result } from "../../dist/index.js";
+import { AttachmentRegistry, nornir, Result } from "../../dist/index.js";
 
 describe("Nornir Tests", () => {
   describe("use()", () => {
     it("Should add basic middleware", async () => {
-      const middlewareMock = jest.fn((input: { test: string; other: true }) => input.other);
-      const chain: (input: { test: string; other: true }) => Promise<boolean> = nornir<{ test: string; other: true }>()
+      const middlewareMock = jest.fn((input: {
+        test: string;
+        other: true;
+      }) => input.other);
+      const chain: (input: {
+        test: string;
+        other: true;
+      }) => Promise<boolean> = nornir<{
+        test: string;
+        other: true;
+      }>()
         .use(middlewareMock)
         .build();
 
-      await expect(chain({ test: "test", other: true })).resolves.toBe(true);
+      await expect(chain({
+        test: "test",
+        other: true,
+      })).resolves.toBe(true);
 
       expect(middlewareMock).toHaveBeenCalledTimes(1);
-      expect(middlewareMock).toHaveBeenCalledWith({ test: "test", other: true }, expect.anything());
+      expect(middlewareMock).toHaveBeenCalledWith({
+        test: "test",
+        other: true,
+      }, expect.anything());
     });
 
     it("Should throw when middleware throws", async () => {
       const middlewareMock = jest.fn(() => {
         throw new Error("boom");
       });
-      const chain: (input: { test: string; other: true }) => Promise<boolean> = nornir<{ test: string; other: true }>()
+      const chain: (input: {
+        test: string;
+        other: true;
+      }) => Promise<boolean> = nornir<{
+        test: string;
+        other: true;
+      }>()
         .use(middlewareMock)
         .build();
 
-      await expect(chain({ test: "test", other: true })).rejects.toThrow("boom");
+      await expect(chain({
+        test: "test",
+        other: true,
+      })).rejects.toThrow("boom");
 
       expect(middlewareMock).toHaveBeenCalledTimes(1);
     });
@@ -31,40 +55,143 @@ describe("Nornir Tests", () => {
 
   describe("useResult()", () => {
     it("Should add basic middleware", async () => {
-      const middlewareMock = jest.fn((input: Result<{ test: string; other: true }>) => input.unwrap().other);
-      const chain: (input: { test: string; other: true }) => Promise<boolean> = nornir<{ test: string; other: true }>()
+      const middlewareMock = jest.fn((
+        input: Result<{
+          test: string;
+          other: true;
+        }>,
+      ) => input.unwrap().other);
+      const chain: (input: {
+        test: string;
+        other: true;
+      }) => Promise<boolean> = nornir<{
+        test: string;
+        other: true;
+      }>()
         .useResult(middlewareMock)
         .build();
 
-      await expect(chain({ test: "test", other: true })).resolves.toBe(true);
+      await expect(chain({
+        test: "test",
+        other: true,
+      })).resolves.toBe(true);
 
       expect(middlewareMock).toHaveBeenCalledTimes(1);
-      expect(middlewareMock).toHaveBeenCalledWith(Result.ok({ test: "test", other: true }), expect.anything());
+      expect(middlewareMock).toHaveBeenCalledWith(
+        Result.ok({
+          test: "test",
+          other: true,
+        }),
+        expect.anything(),
+      );
     });
 
     it("Should handle previous error result", async () => {
-      const middlewareMock = jest.fn((result: Result<{ test: string; other: true }>) =>
+      const middlewareMock = jest.fn((
+        result: Result<{
+          test: string;
+          other: true;
+        }>,
+      ) =>
         result
           .chain(() => Result.ok(false), () => Result.ok(true)).unwrap()
       );
-      const chain: (input: { test: string; other: true }) => Promise<boolean> = nornir<{ test: string; other: true }>()
+      const chain: (input: {
+        test: string;
+        other: true;
+      }) => Promise<boolean> = nornir<{
+        test: string;
+        other: true;
+      }>()
         .use(() => {
           throw new Error("boom");
         })
         .useResult(middlewareMock)
         .build();
 
-      await expect(chain({ test: "test", other: true })).resolves.toBe(true);
+      await expect(chain({
+        test: "test",
+        other: true,
+      })).resolves.toBe(true);
       expect(middlewareMock).toHaveBeenCalledTimes(1);
       expect(middlewareMock).toHaveBeenCalledWith(Result.err(new Error("boom")), expect.anything());
     });
   });
 
+  describe("build() with base registry", () => {
+    it("Should build a chain a registry including existing registry keys", async () => {
+      const registry = new AttachmentRegistry();
+      const factory = jest.fn(() => "other");
+      const key = registry.registerFactory(factory);
+
+      const middlewareMock = jest.fn((input: {
+        test: string;
+        other: true;
+      }, registry: AttachmentRegistry) => {
+        return registry.get(key);
+      });
+      const chain = nornir<{
+        test: string;
+        other: true;
+      }>()
+        .use(middlewareMock)
+        .build(registry);
+
+      await expect(chain({
+        test: "test",
+        other: true,
+      })).resolves.toBe("other");
+
+      expect(factory).toHaveBeenCalledTimes(1);
+      expect(middlewareMock).toHaveBeenCalledTimes(1);
+      expect(middlewareMock).toHaveBeenCalledWith({
+        test: "test",
+        other: true,
+      }, expect.anything());
+    });
+
+    it("Base registry should not be modified", async () => {
+      const registry = new AttachmentRegistry();
+      const factory = jest.fn(() => "other");
+      const key = registry.registerFactory(factory);
+      const modKey = AttachmentRegistry.createKey<string>();
+
+      const middlewareMock = jest.fn((input: {
+        test: string;
+        other: true;
+      }, middlewareRegistry: AttachmentRegistry) => {
+        middlewareRegistry.put(modKey, "other");
+        return middlewareRegistry.get(key);
+      });
+      const chain = nornir<{
+        test: string;
+        other: true;
+      }>()
+        .use(middlewareMock)
+        .build(registry);
+
+      await chain({
+        test: "test",
+        other: true,
+      });
+
+      expect(factory).toHaveBeenCalledTimes(1);
+      expect(registry.get(key)).toBe("other");
+      expect(registry.get(modKey)).toBeNull();
+    });
+  });
+
   describe("useChain()", () => {
     it("Should compose chain", async () => {
       const nestedMiddlewareMock = jest.fn((input: string) => input.length);
       const middlewareMock = jest.fn((input: number) => input === 4);
-      const chain: (input: { test: string; other: true }) => Promise<boolean> = nornir<{ test: string; other: true }>()
+      const chain: (input: {
+        test: string;
+        other: true;
+      }) => Promise<boolean> = nornir<{
+        test: string;
+        other: true;
+      }>()
         .use(input => input.test)
         .useChain(
           nornir<string>()
@@ -73,7 +200,10 @@ describe("Nornir Tests", () => {
         .use(middlewareMock)
         .build();
 
-      await expect(chain({ test: "test", other: true })).resolves.toBe(true);
+      await expect(chain({
+        test: "test",
+        other: true,
+      })).resolves.toBe(true);
       expect(nestedMiddlewareMock).toHaveBeenCalledTimes(1);
       expect(nestedMiddlewareMock).toHaveBeenCalledWith("test", expect.anything());
       expect(middlewareMock).toHaveBeenCalledTimes(1);
@@ -83,28 +213,54 @@ describe("Nornir Tests", () => {
 
   describe("split()", () => {
     it("Should split items", async () => {
-      const splitMiddlewareMock = jest.fn((input: { test: string; other: true }) => input.test);
+      const splitMiddlewareMock = jest.fn((input: {
+        test: string;
+        other: true;
+      }) => input.test);
       const middlewareMock = jest.fn((items: Result<string, Error>[]) => items.map(item => item.unwrap()));
-      const chain: (input: { test: string; other: true }[]) => Promise<string[]> = nornir<
-        { test: string; other: true }[]
+      const chain: (input: {
+        test: string;
+        other: true;
+      }[]) => Promise<string[]> = nornir<
+        {
+          test: string;
+          other: true;
+        }[]
       >()
         .split(nested => nested.use(splitMiddlewareMock))
         .use(middlewareMock)
         .build();
 
-      await expect(chain([{ test: "test", other: true }, { test: "test2", other: true }])).resolves.toEqual([
+      await expect(chain([{
+        test: "test",
+        other: true,
+      }, {
+        test: "test2",
+        other: true,
+      }])).resolves.toEqual([
         "test",
         "test2",
       ]);
       expect(splitMiddlewareMock).toHaveBeenCalledTimes(2);
-      expect(splitMiddlewareMock).toHaveBeenCalledWith({ test: "test", other: true }, expect.anything());
-      expect(splitMiddlewareMock).toHaveBeenCalledWith({ test: "test2", other: true }, expect.anything());
+      expect(splitMiddlewareMock).toHaveBeenCalledWith({
+        test: "test",
+        other: true,
+      }, expect.anything());
+      expect(splitMiddlewareMock).toHaveBeenCalledWith({
+        test: "test2",
+        other: true,
+      }, expect.anything());
       expect(middlewareMock).toHaveBeenCalledTimes(1);
       expect(middlewareMock).toHaveBeenCalledWith([Result.ok("test"), Result.ok("test2")], expect.anything());
     });
 
     it("Should handle errors", async () => {
-      const splitMiddlewareMock = jest.fn<(input: { test: string; other: true }) => string>();
+      const splitMiddlewareMock = jest.fn<
+        (input: {
+          test: string;
+          other: true;
+        }) => string
+      >();
       splitMiddlewareMock
         .mockImplementationOnce(input => input.test)
         .mockRejectedValueOnce(new Error("boom") as never);
@@ -112,20 +268,38 @@ describe("Nornir Tests", () => {
       const middlewareMock = jest.fn((items: Result<string, Error>[]) =>
         items.map(item => item.isOk ? item.unwrap() : item.error.message)
       );
-      const chain: (input: { test: string; other: true }[]) => Promise<string[]> = nornir<
-        { test: string; other: true }[]
+      const chain: (input: {
+        test: string;
+        other: true;
+      }[]) => Promise<string[]> = nornir<
+        {
+          test: string;
+          other: true;
+        }[]
       >()
         .split(nested => nested.use(splitMiddlewareMock))
         .use(middlewareMock)
         .build();
 
-      await expect(chain([{ test: "test", other: true }, { test: "test2", other: true }])).resolves.toEqual([
+      await expect(chain([{
+        test: "test",
+        other: true,
+      }, {
+        test: "test2",
+        other: true,
+      }])).resolves.toEqual([
         "test",
         "boom",
       ]);
       expect(splitMiddlewareMock).toHaveBeenCalledTimes(2);
-      expect(splitMiddlewareMock).toHaveBeenCalledWith({ test: "test", other: true }, expect.anything());
-      expect(splitMiddlewareMock).toHaveBeenCalledWith({ test: "test2", other: true }, expect.anything());
+      expect(splitMiddlewareMock).toHaveBeenCalledWith({
+        test: "test",
+        other: true,
+      }, expect.anything());
+      expect(splitMiddlewareMock).toHaveBeenCalledWith({
+        test: "test2",
+        other: true,
+      }, expect.anything());
       expect(middlewareMock).toHaveBeenCalledTimes(1);
       expect(middlewareMock).toHaveBeenCalledWith(
         [Result.ok("test"), Result.err(new Error("boom"))],
@@ -158,9 +332,15 @@ describe("Nornir Tests", () => {
         })
         .build();
 
-      await expect(chain({ type: "test1", value: "test" })).resolves.toBe("test");
+      await expect(chain({
+        type: "test1",
+        value: "test",
+      })).resolves.toBe("test");
       expect(test1MiddlewareMock).toHaveBeenCalledTimes(1);
-      expect(test1MiddlewareMock).toHaveBeenCalledWith({ type: "test1", value: "test" }, expect.anything());
+      expect(test1MiddlewareMock).toHaveBeenCalledWith({
+        type: "test1",
+        value: "test",
+      }, expect.anything());
       expect(test2MiddlewareMock).not.toHaveBeenCalled();
     });
   });
diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts
index 9f264cd..ab59e18 100644
--- a/packages/core/src/index.ts
+++ b/packages/core/src/index.ts
@@ -1,4 +1,4 @@
-export { AttachmentKey, AttachmentRegistry } from "./lib/attachment-registry.js";
+export { AttachmentKey, AttachmentRegistry, RegistryFactory } from "./lib/attachment-registry.js";
 export { NornirBuildError, NornirMissingAttachmentException, NornirValidationError } from "./lib/error.js";
 export { InitialArgumentsKey, Nornir, nornir, nornir as default } from "./lib/nornir.js";
 export { Result } from "./lib/result.js";
diff --git a/packages/core/src/lib/attachment-registry.ts b/packages/core/src/lib/attachment-registry.ts
index 07a794e..7fe146b 100644
--- a/packages/core/src/lib/attachment-registry.ts
+++ b/packages/core/src/lib/attachment-registry.ts
@@ -5,8 +5,23 @@ export interface AttachmentKey<T> {
   readonly marker: T;
 }
 
+interface ConstantRegistryItem<T = unknown> {
+  type: "constant";
+  value: T;
+}
+
+export type RegistryFactory<T> = (registry: AttachmentRegistry) => T;
+
+interface FactoryRegistryItem<T = unknown> {
+  type: "factory";
+  factory: RegistryFactory<T>;
+  value?: T;
+}
+
+type RegistryItem<T = unknown> = ConstantRegistryItem<T> | FactoryRegistryItem<T>;
+
 export class AttachmentRegistry {
-  private attachments = new Map<symbol, unknown>();
+  private attachments = new Map<symbol, RegistryItem>();
 
   public static createKey<T>(): AttachmentKey<T> {
     return {
@@ -15,8 +30,30 @@ export class AttachmentRegistry {
     };
   }
 
+  public register<T>(value: T) {
+    const key = AttachmentRegistry.createKey<T>();
+    this.put(key, value);
+    return key;
+  }
+
+  public registerFactory<T>(factory: RegistryFactory<T>) {
+    const key = AttachmentRegistry.createKey<T>();
+    this.putFactory(key, factory);
+    return key;
+  }
+
   public put<T>(key: AttachmentKey<T>, value: T): void {
-    this.attachments.set(key.id, value);
+    this.attachments.set(key.id, {
+      type: "constant",
+      value,
+    });
+  }
+
+  public putFactory<T>(key: AttachmentKey<T>, factory: RegistryFactory<T>): void {
+    this.attachments.set(key.id, {
+      type: "factory",
+      factory,
+    });
   }
 
   public delete<T>(key: AttachmentKey<T>): void {
@@ -24,14 +61,27 @@ export class AttachmentRegistry {
   }
 
   public get<T>(key: AttachmentKey<T>): T | null {
-    return this.attachments.get(key.id) as T;
+    if (this.has(key)) {
+      return this.resolveItem(this.attachments.get(key.id) as RegistryItem<T>);
+    }
+    return null;
+  }
+
+  private resolveItem<T>(item: RegistryItem<T>): T {
+    if (item.type === "constant") {
+      return item.value;
+    }
+    if (item.value === undefined) {
+      item.value = item.factory(this);
+    }
+    return item.value;
   }
 
   public getAssert<T>(key: AttachmentKey<T>): T {
     if (!this.has(key)) {
       throw new NornirMissingAttachmentException();
     }
-    return this.attachments.get(key.id) as T;
+    return this.resolveItem(this.attachments.get(key.id)! as RegistryItem<T>);
   }
 
   public has<T>(key: AttachmentKey<T>): boolean {
@@ -57,4 +107,12 @@ export class AttachmentRegistry {
     }
     return false;
   }
+
+  public merge(...registries: AttachmentRegistry[]) {
+    for (const registry of registries) {
+      for (const [key, value] of registry.attachments) {
+        this.attachments.set(key, value);
+      }
+    }
+  }
 }
diff --git a/packages/core/src/lib/nornir.ts b/packages/core/src/lib/nornir.ts
index ff80c45..a9c0c16 100644
--- a/packages/core/src/lib/nornir.ts
+++ b/packages/core/src/lib/nornir.ts
@@ -17,9 +17,12 @@ class NornirContext {
     this.chain.push(handler);
   }
 
-  public build() {
+  public build(baseRegistry?: AttachmentRegistry) {
     return async (...args: unknown[]) => {
       const registry = new AttachmentRegistry();
+      if (baseRegistry) {
+        registry.merge(baseRegistry);
+      }
       let input = Result.ok(args[0]);
       if (!registry.has(InitialArgumentsKey)) {
         registry.put(InitialArgumentsKey, args);
@@ -177,8 +180,8 @@ export class Nornir<Input, StepInput = Input> {
   /**
    * Build the chain into a function that takes the chain input and produces the chain output.
    */
-  public build() {
-    return this.context.build() as (input: Input) => Promise<StepInput>;
+  public build(baseRegistry?: AttachmentRegistry) {
+    return this.context.build(baseRegistry) as (input: Input) => Promise<StepInput>;
   }
 
   public buildWithContext() {