diff --git a/packages/durabletask-js-azuremanaged/src/worker-builder.ts b/packages/durabletask-js-azuremanaged/src/worker-builder.ts index bffd09fe..f749ea79 100644 --- a/packages/durabletask-js-azuremanaged/src/worker-builder.ts +++ b/packages/durabletask-js-azuremanaged/src/worker-builder.ts @@ -10,6 +10,7 @@ import { TActivity, TInput, TOutput, + EntityFactory, Logger, ConsoleLogger, VersioningOptions, @@ -25,6 +26,7 @@ export class DurableTaskAzureManagedWorkerBuilder { private _grpcChannelOptions: grpc.ChannelOptions = {}; private _orchestrators: { name?: string; fn: TOrchestrator }[] = []; private _activities: { name?: string; fn: TActivity }[] = []; + private _entities: { name?: string; factory: EntityFactory }[] = []; private _logger: Logger = new ConsoleLogger(); private _shutdownTimeoutMs?: number; private _versioning?: VersioningOptions; @@ -185,6 +187,30 @@ export class DurableTaskAzureManagedWorkerBuilder { return this; } + /** + * Registers an entity factory with the worker. + * The entity name is derived from the factory function name. + * + * @param factory The entity factory function. + * @returns This builder instance. + */ + addEntity(factory: EntityFactory): DurableTaskAzureManagedWorkerBuilder { + this._entities.push({ factory }); + return this; + } + + /** + * Registers a named entity factory with the worker. + * + * @param name The name of the entity. + * @param factory The entity factory function. + * @returns This builder instance. + */ + addNamedEntity(name: string, factory: EntityFactory): DurableTaskAzureManagedWorkerBuilder { + this._entities.push({ name, factory }); + return this; + } + /** * Sets the logger to use for logging. * Defaults to ConsoleLogger. @@ -289,6 +315,15 @@ export class DurableTaskAzureManagedWorkerBuilder { } } + // Register all entities + for (const { name, factory } of this._entities) { + if (name) { + worker.addNamedEntity(name, factory); + } else { + worker.addEntity(factory); + } + } + return worker; } } diff --git a/packages/durabletask-js-azuremanaged/test/unit/worker-builder.spec.ts b/packages/durabletask-js-azuremanaged/test/unit/worker-builder.spec.ts new file mode 100644 index 00000000..d0f97cb4 --- /dev/null +++ b/packages/durabletask-js-azuremanaged/test/unit/worker-builder.spec.ts @@ -0,0 +1,93 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +import { DurableTaskAzureManagedWorkerBuilder, createAzureManagedWorkerBuilder } from "../../src/worker-builder"; +import { TaskEntity, ITaskEntity, TaskEntityOperation } from "@microsoft/durabletask-js"; + +// Simple test entity for registration testing +class CounterEntity extends TaskEntity { + add(operation: TaskEntityOperation): void { + const amount = operation.getInput() ?? 1; + this.state = (this.state ?? 0) + amount; + } +} + +// Factory functions for testing +function createCounterEntity(): ITaskEntity { + return new CounterEntity(); +} + +describe("DurableTaskAzureManagedWorkerBuilder", () => { + const ENDPOINT = "localhost:8080"; + const TASKHUB = "test"; + + describe("addEntity", () => { + it("should register an entity factory and return the builder for chaining", () => { + const builder = new DurableTaskAzureManagedWorkerBuilder(); + + const result = builder.endpoint(ENDPOINT, TASKHUB, null).addEntity(createCounterEntity); + + expect(result).toBe(builder); + }); + + it("should register an entity that gets added to the worker on build", () => { + const builder = new DurableTaskAzureManagedWorkerBuilder(); + + const worker = builder.endpoint(ENDPOINT, TASKHUB, null).addEntity(createCounterEntity).build(); + + // The worker should have the entity registered. We verify by checking that + // attempting to register it again with the same name throws a duplicate error. + expect(() => worker.addNamedEntity("createCounterEntity", createCounterEntity)).toThrow(); + }); + }); + + describe("addNamedEntity", () => { + it("should register a named entity factory and return the builder for chaining", () => { + const builder = new DurableTaskAzureManagedWorkerBuilder(); + + const result = builder.endpoint(ENDPOINT, TASKHUB, null).addNamedEntity("MyCounter", createCounterEntity); + + expect(result).toBe(builder); + }); + + it("should register a named entity that gets added to the worker on build", () => { + const builder = new DurableTaskAzureManagedWorkerBuilder(); + + const worker = builder.endpoint(ENDPOINT, TASKHUB, null).addNamedEntity("MyCounter", createCounterEntity).build(); + + // Entity names are lowercased. Registering the same lowercased name should throw. + expect(() => worker.addNamedEntity("mycounter", createCounterEntity)).toThrow(); + }); + }); + + describe("fluent chaining with entities", () => { + it("should support registering orchestrators, activities, and entities together", () => { + const orchestrator = async function* testOrchestrator() { + yield; + return "done"; + }; + const activity = async () => "result"; + + const builder = new DurableTaskAzureManagedWorkerBuilder(); + + const worker = builder + .endpoint(ENDPOINT, TASKHUB, null) + .addNamedOrchestrator("myOrch", orchestrator) + .addNamedActivity("myActivity", activity) + .addNamedEntity("MyCounter", createCounterEntity) + .build(); + + expect(worker).toBeDefined(); + }); + }); + + describe("createAzureManagedWorkerBuilder", () => { + it("should create a builder that supports entity registration", () => { + const builder = createAzureManagedWorkerBuilder(ENDPOINT, TASKHUB, null); + + const result = builder.addNamedEntity("MyCounter", createCounterEntity); + + expect(result).toBe(builder); + }); + }); +});