in-memory-backend.spec.ts

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

import {
  InMemoryOrchestrationBackend,
  TestOrchestrationClient,
  TestOrchestrationWorker,
  OrchestrationStatus,
  getName,
  whenAll,
  ActivityContext,
  OrchestrationContext,
  Task,
  TOrchestrator,
} from "../src";

describe("In-Memory Backend", () => {
  let backend: InMemoryOrchestrationBackend;
  let client: TestOrchestrationClient;
  let worker: TestOrchestrationWorker;

  beforeEach(async () => {
    backend = new InMemoryOrchestrationBackend();
    client = new TestOrchestrationClient(backend);
    worker = new TestOrchestrationWorker(backend);
  });

  afterEach(async () => {
    if (worker) {
      try {
        await worker.stop();
      } catch {
        // Ignore if not running
      }
    }
    backend.reset();
  });

  it("should run an empty orchestration", async () => {
    let invoked = false;

    const emptyOrchestrator: TOrchestrator = async (_: OrchestrationContext) => {
      invoked = true;
    };

    worker.addOrchestrator(emptyOrchestrator);
    await worker.start();

    const id = await client.scheduleNewOrchestration(emptyOrchestrator);
    const state = await client.waitForOrchestrationCompletion(id, true, 10);

    expect(invoked).toBe(true);
    expect(state).toBeDefined();
    expect(state?.name).toEqual(getName(emptyOrchestrator));
    expect(state?.instanceId).toEqual(id);
    expect(state?.failureDetails).toBeUndefined();
    expect(state?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED);
  });

  it("should run an activity sequence", async () => {
    const plusOne = async (_: ActivityContext, input: number) => {
      return input + 1;
    };

    const sequence: TOrchestrator = async function* (ctx: OrchestrationContext, startVal: number): any {
      const numbers = [startVal];
      let current = startVal;

      for (let i = 0; i < 5; i++) {
        current = yield ctx.callActivity(plusOne, current);
        numbers.push(current);
      }

      return numbers;
    };

    worker.addOrchestrator(sequence);
    worker.addActivity(plusOne);
    await worker.start();

    const id = await client.scheduleNewOrchestration(sequence, 1);
    const state = await client.waitForOrchestrationCompletion(id, true, 10);

    expect(state).toBeDefined();
    expect(state?.name).toEqual(getName(sequence));
    expect(state?.failureDetails).toBeUndefined();
    expect(state?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED);
    expect(state?.serializedInput).toEqual(JSON.stringify(1));
    expect(state?.serializedOutput).toEqual(JSON.stringify([1, 2, 3, 4, 5, 6]));
  });

  it("should run fan-out/fan-in", async () => {
    let activityCounter = 0;

    const increment = (_: ActivityContext) => {
      activityCounter++;
    };

    const orchestrator: TOrchestrator = async function* (ctx: OrchestrationContext, count: number): any {
      const tasks: Task<any>[] = [];

      for (let i = 0; i < count; i++) {
        tasks.push(ctx.callActivity(increment));
      }

      yield whenAll(tasks);
    };

    worker.addActivity(increment);
    worker.addOrchestrator(orchestrator);
    await worker.start();

    const id = await client.scheduleNewOrchestration(orchestrator, 5);
    const state = await client.waitForOrchestrationCompletion(id, true, 10);

    expect(state).toBeDefined();
    expect(state?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED);
    expect(activityCounter).toEqual(5);
  });

  it("should handle sub-orchestrations", async () => {
    let activityCounter = 0;

    const increment = (_: ActivityContext) => {
      activityCounter++;
    };

    const childOrchestrator: TOrchestrator = async function* (ctx: OrchestrationContext): any {
      yield ctx.callActivity(increment);
    };

    const parentOrchestrator: TOrchestrator = async function* (ctx: OrchestrationContext): any {
      yield ctx.callSubOrchestrator(childOrchestrator);
    };

    worker.addActivity(increment);
    worker.addOrchestrator(childOrchestrator);
    worker.addOrchestrator(parentOrchestrator);
    await worker.start();

    const id = await client.scheduleNewOrchestration(parentOrchestrator);
    const state = await client.waitForOrchestrationCompletion(id, true, 10);

    expect(state).toBeDefined();
    expect(state?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED);
    expect(activityCounter).toEqual(1);
  });

  it("should handle external events", async () => {
    const orchestrator: TOrchestrator = async function* (ctx: OrchestrationContext): any {
      const value = yield ctx.waitForExternalEvent("my_event");
      return value;
    };

    worker.addOrchestrator(orchestrator);
    await worker.start();

    const id = await client.scheduleNewOrchestration(orchestrator);
    
    // Wait for orchestration to start
    await client.waitForOrchestrationStart(id, false, 5);
    
    // Raise the event
    await client.raiseOrchestrationEvent(id, "my_event", "hello");
    
    const state = await client.waitForOrchestrationCompletion(id, true, 10);

    expect(state).toBeDefined();
    expect(state?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED);
    expect(state?.serializedOutput).toEqual(JSON.stringify("hello"));
  });

  it("should handle timers", async () => {
    const orchestrator: TOrchestrator = async function* (ctx: OrchestrationContext): any {
      // Wait for 100ms
      yield ctx.createTimer(0.1);
      return "done";
    };

    worker.addOrchestrator(orchestrator);
    await worker.start();

    const id = await client.scheduleNewOrchestration(orchestrator);
    const state = await client.waitForOrchestrationCompletion(id, true, 10);

    expect(state).toBeDefined();
    expect(state?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED);
    expect(state?.serializedOutput).toEqual(JSON.stringify("done"));
  });

  it("should handle termination", async () => {
    const orchestrator: TOrchestrator = async function* (ctx: OrchestrationContext): any {
      yield ctx.waitForExternalEvent("never");
      return "never reached";
    };

    worker.addOrchestrator(orchestrator);
    await worker.start();

    const id = await client.scheduleNewOrchestration(orchestrator);
    await client.waitForOrchestrationStart(id, false, 5);
    
    await client.terminateOrchestration(id, "terminated by test");
    
    const state = await client.waitForOrchestrationCompletion(id, true, 10);

    expect(state).toBeDefined();
    expect(state?.runtimeStatus).toEqual(OrchestrationStatus.TERMINATED);
    expect(state?.serializedOutput).toEqual(JSON.stringify("terminated by test"));
  });

  it("should handle continue-as-new", async () => {
    const orchestrator: TOrchestrator = async (ctx: OrchestrationContext, input: number) => {
      if (input < 5) {
        ctx.continueAsNew(input + 1, true);
      } else {
        return input;
      }
    };

    worker.addOrchestrator(orchestrator);
    await worker.start();

    const id = await client.scheduleNewOrchestration(orchestrator, 1);
    const state = await client.waitForOrchestrationCompletion(id, true, 10);

    expect(state).toBeDefined();
    expect(state?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED);
    expect(state?.serializedOutput).toEqual(JSON.stringify(5));
  });

  it("should clear customStatus after continue-as-new", async () => {
    const orchestrator: TOrchestrator = async (ctx: OrchestrationContext, input: number) => {
      if (input === 1) {
        // First iteration: set a custom status then continue-as-new
        ctx.setCustomStatus("iteration-1-status");
        ctx.continueAsNew(2, false);
      } else {
        // Second iteration: do NOT set custom status — it should be cleared
        return "done";
      }
    };

    worker.addOrchestrator(orchestrator);
    await worker.start();

    const id = await client.scheduleNewOrchestration(orchestrator, 1);
    const state = await client.waitForOrchestrationCompletion(id, true, 10);

    expect(state).toBeDefined();
    expect(state?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED);
    expect(state?.serializedOutput).toEqual(JSON.stringify("done"));
    // customStatus must be cleared after continue-as-new when the new iteration
    // does not set one — it should not carry over from the previous iteration
    expect(state?.serializedCustomStatus).toBeUndefined();
  });

  it("should preserve sendEvent actions when continuing-as-new", async () => {
    // Receiver orchestration that waits for an event
    const receiver: TOrchestrator = async function* (ctx: OrchestrationContext): any {
      const value = yield ctx.waitForExternalEvent("ping");
      return value;
    };

    // Sender orchestration that sends an event then continues-as-new
    const sender: TOrchestrator = async (ctx: OrchestrationContext, input: { receiverId: string; iteration: number }) => {
      if (input.iteration === 1) {
        // On first iteration, send event to receiver then continue-as-new
        ctx.sendEvent(input.receiverId, "ping", "hello from sender");
        ctx.continueAsNew({ receiverId: input.receiverId, iteration: 2 }, false);
      } else {
        return "sender done";
      }
    };

    worker.addOrchestrator(receiver);
    worker.addOrchestrator(sender);
    await worker.start();

    // Start receiver first, then sender
    const receiverId = await client.scheduleNewOrchestration(receiver);
    await client.waitForOrchestrationStart(receiverId, false, 5);

    const senderId = await client.scheduleNewOrchestration(sender, { receiverId, iteration: 1 });

    // Wait for both to complete
    const senderState = await client.waitForOrchestrationCompletion(senderId, true, 10);
    const receiverState = await client.waitForOrchestrationCompletion(receiverId, true, 10);

    // Sender should complete after continuing-as-new
    expect(senderState).toBeDefined();
    expect(senderState?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED);
    expect(senderState?.serializedOutput).toEqual(JSON.stringify("sender done"));

    // Receiver should have received the event sent before continue-as-new
    expect(receiverState).toBeDefined();
    expect(receiverState?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED);
    expect(receiverState?.serializedOutput).toEqual(JSON.stringify("hello from sender"));
  });

  it("should handle orchestration without activities", async () => {
    const orchestrator: TOrchestrator = async (_: OrchestrationContext, input: number) => {
      return input * 2;
    };

    worker.addOrchestrator(orchestrator);
    await worker.start();

    const id = await client.scheduleNewOrchestration(orchestrator, 21);
    const state = await client.waitForOrchestrationCompletion(id, true, 10);

    expect(state).toBeDefined();
    expect(state?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED);
    expect(state?.serializedOutput).toEqual(JSON.stringify(42));
  });

  it("should handle activity failures", async () => {
    const failingActivity = (_: ActivityContext) => {
      throw new Error("Activity failed intentionally");
    };

    const orchestrator: TOrchestrator = async function* (ctx: OrchestrationContext): any {
      try {
        yield ctx.callActivity(failingActivity);
        return "should not reach here";
      } catch (error: any) {
        return `caught: ${error.message}`;
      }
    };

    worker.addActivity(failingActivity);
    worker.addOrchestrator(orchestrator);
    await worker.start();

    const id = await client.scheduleNewOrchestration(orchestrator);
    const state = await client.waitForOrchestrationCompletion(id, true, 10);

    expect(state).toBeDefined();
    expect(state?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED);
    expect(state?.serializedOutput).toContain("caught:");
  });

  it("should purge completed orchestrations", async () => {
    const orchestrator: TOrchestrator = async () => "done";

    worker.addOrchestrator(orchestrator);
    await worker.start();

    const id = await client.scheduleNewOrchestration(orchestrator);
    await client.waitForOrchestrationCompletion(id, false, 10);

    const result = await client.purgeOrchestration(id);
    expect(result.deletedInstanceCount).toEqual(1);

    const state = await client.getOrchestrationState(id);
    expect(state).toBeUndefined();
  });

  it("should continue processing after backend reset during orchestration execution", async () => {
    // This test verifies that the worker's processing loop survives a backend
    // reset that occurs while an orchestration is being processed. When the
    // instance is deleted (e.g., via reset) before completeOrchestration runs,
    // the worker must not crash. Without proper handling, the processing loop
    // would terminate, preventing any subsequent orchestrations from running.

    const validOrchestrator: TOrchestrator = async (_: OrchestrationContext, input: number) => {
      return input + 1;
    };

    worker.addOrchestrator(validOrchestrator);
    await worker.start();

    // Schedule and complete a first orchestration to verify baseline behavior
    const id1 = await client.scheduleNewOrchestration(validOrchestrator, 10);
    const state1 = await client.waitForOrchestrationCompletion(id1, true, 10);
    expect(state1?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED);
    expect(state1?.serializedOutput).toEqual(JSON.stringify(11));

    // Reset the backend (simulates clearing all state while worker is running)
    backend.reset();

    // Schedule a second orchestration after reset — this verifies the worker
    // is still alive and can process new work items
    const id2 = await client.scheduleNewOrchestration(validOrchestrator, 41);
    const state2 = await client.waitForOrchestrationCompletion(id2, true, 10);

    expect(state2).toBeDefined();
    expect(state2?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED);
    expect(state2?.serializedOutput).toEqual(JSON.stringify(42));
  });

  it("should silently ignore completeOrchestration for purged instances", () => {
    // Verifies that completeOrchestration returns silently when the instance
    // has been deleted (e.g., via purge or reset), consistent with how
    // completeActivity handles missing instances.

    // Should not throw — instance simply doesn't exist
    expect(() => {
      backend.completeOrchestration("nonexistent-instance", 1, []);
    }).not.toThrow();
  });

  it("should allow reusing instance IDs after reset", async () => {
    const orchestrator: TOrchestrator = async (_: OrchestrationContext, input: number) => {
      return input * 2;
    };

    // Create an orchestration without starting the worker, so it stays in the queue
    const instanceId = "reuse-test-id";
    backend.createInstance(instanceId, getName(orchestrator), JSON.stringify(10));

    // Reset while the orchestration is still queued (not yet processed)
    backend.reset();

    // Now create a new orchestration with the same instance ID and process it
    worker.addOrchestrator(orchestrator);
    await worker.start();

    await client.scheduleNewOrchestration(orchestrator, 21, instanceId);
    const state = await client.waitForOrchestrationCompletion(instanceId, true, 10);
    expect(state?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED);
    expect(state?.serializedOutput).toEqual(JSON.stringify(42));
  });
});