Fix service handler options propagation (#405)

slinkydeveloper · web-flow · commit f89e2f7c9aba · 2024-07-16T14:29:41.000+02:00
* Patch wrong propagation of ServiceHandlerOpts

* Add test for the service handler with raw input.

Also refactored the internal unit test infra to decouple the execution of the state machine from the service definition.
diff --git a/packages/restate-sdk/src/types/rpc.ts b/packages/restate-sdk/src/types/rpc.ts
@@ -438,10 +438,10 @@ export const service = <P extends string, M>(service: {
     throw new Error("service must be defined");
   }
   const handlers = Object.entries(service.handlers).map(([name, handler]) => {
-    if (handler instanceof HandlerWrapper) {
-      return [name, handler.transpose()];
-    }
     if (handler instanceof Function) {
+      if (HandlerWrapper.fromHandler(handler) !== undefined) {
+        return [name, handler];
+      }
       return [
         name,
         HandlerWrapper.from(HandlerKind.SERVICE, handler).transpose(),
diff --git a/packages/restate-sdk/test/service_bind.test.ts b/packages/restate-sdk/test/service_bind.test.ts
@@ -10,10 +10,16 @@
  */
 
 import type { TestGreeter, TestRequest } from "./testdriver.js";
-import { TestDriver, TestResponse } from "./testdriver.js";
-import type * as restate from "../src/public_api.js";
-import { greetRequest, inputMessage, startMessage } from "./protoutils.js";
-import { describe, it } from "vitest";
+import { TestDriver, TestResponse, testService } from "./testdriver.js";
+import * as restate from "../src/public_api.js";
+import {
+  END_MESSAGE,
+  greetRequest,
+  inputMessage,
+  outputMessage,
+  startMessage,
+} from "./protoutils.js";
+import { describe, expect, it } from "vitest";
 
 const greeter: TestGreeter = {
   // eslint-disable-next-line @typescript-eslint/require-await
@@ -57,3 +63,32 @@ describe("BindService", () => {
     ]).run();
   });
 });
+
+const acceptBytes = restate.service({
+  name: "acceptBytes",
+  handlers: {
+    greeter: restate.handlers.handler(
+      {
+        accept: "application/octet-stream",
+        contentType: "application/json",
+      },
+      // eslint-disable-next-line @typescript-eslint/require-await
+      async (_ctx: restate.Context, audio: Uint8Array) => {
+        return { length: audio.length };
+      }
+    ),
+  },
+});
+
+describe("AcceptBytes", () => {
+  it("should accept bytes", async () => {
+    const result = await testService(acceptBytes).run({
+      input: [startMessage(), inputMessage(new Uint8Array([0, 1, 2, 3, 4]))],
+    });
+
+    expect(result).toStrictEqual([
+      outputMessage(Buffer.from(JSON.stringify({ length: 5 }))),
+      END_MESSAGE,
+    ]);
+  });
+});
diff --git a/packages/restate-sdk/test/testdriver.ts b/packages/restate-sdk/test/testdriver.ts
@@ -16,14 +16,16 @@ import {
   StartMessage,
 } from "../src/types/protocol.js";
 import type { Connection } from "../src/connection/connection.js";
-import { formatMessageAsJson } from "../src/utils/utils.js";
 import { Message } from "../src/types/types.js";
 import { StateMachine } from "../src/state_machine.js";
 import { InvocationBuilder } from "../src/invocation.js";
 import type { ObjectContext } from "../src/context.js";
-import type { VirtualObjectDefinition } from "../src/public_api.js";
+import type {
+  ServiceDefinition,
+  VirtualObjectDefinition,
+  WorkflowDefinition,
+} from "../src/public_api.js";
 import { object } from "../src/public_api.js";
-import { HandlerKind } from "../src/types/rpc.js";
 import { NodeEndpoint } from "../src/endpoint/node_endpoint.js";
 import type { EndpointBuilder } from "../src/endpoint/endpoint_builder.js";
 
@@ -47,53 +49,115 @@ export interface TestGreeter {
   greet(ctx: ObjectContext, message: TestRequest): Promise<TestResponse>;
 }
 
-export class TestDriver implements Connection {
-  private readonly result: Message[] = [];
-
-  private restateServer: TestRestateServer;
-  private stateMachine: StateMachine;
-  private completionMessages: Message[];
+export class TestDriver {
+  private readonly uut: UUT<string, unknown>;
+  private readonly input: Message[];
 
+  // Deprecated, please use testService below
   constructor(instance: TestGreeter, entries: Message[]) {
-    this.restateServer = new TestRestateServer();
-
-    const svc = object({
-      name: "greeter",
-      handlers: {
-        greet: async (ctx: ObjectContext, arg: TestRequest) => {
-          return instance.greet(ctx, arg);
+    this.uut = testService(
+      object({
+        name: "greeter",
+        handlers: {
+          greet: async (ctx: ObjectContext, arg: TestRequest) => {
+            return instance.greet(ctx, arg);
+          },
         },
-      },
+      })
+    );
+    this.input = entries;
+  }
+
+  async run(): Promise<Message[]> {
+    return await this.uut.run({
+      input: this.input,
     });
+  }
+}
 
-    this.restateServer.bind(svc);
+/**
+ * This class' only purpose is to make certain methods accessible in tests.
+ * Those methods are otherwise protected, to reduce the public interface and
+ * make it simpler for users to understand what methods are relevant for them,
+ * and which ones are not.
+ */
+class TestRestateServer extends NodeEndpoint {}
+
+interface RunOptions {
+  /// If not provided, will call the first service
+  service?: string;
+  /// If not provided, will call the first handler
+  handler?: string;
+  input: Message[];
+}
+
+export class UUT<N extends string, T> {
+  private readonly defaultService: string;
+  private readonly defaultHandler: string;
+
+  // eslint-disable-next-line @typescript-eslint/no-redundant-type-constituents
+  constructor(
+    private readonly definition:
+      | ServiceDefinition<N, T>
+      | VirtualObjectDefinition<N, T>
+      | WorkflowDefinition<N, T>
+  ) {
+    // Infer service name and handler
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access,@typescript-eslint/no-unsafe-assignment
+    this.defaultService = definition.name;
+
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+    const definitionRecord: Record<string, never> =
+      definition as unknown as Record<string, never>;
+    if (definitionRecord && definitionRecord.service != undefined) {
+      this.defaultHandler = Object.keys(
+        definitionRecord.service as { [s: string]: unknown }
+      )[0];
+    } else if (definitionRecord && definitionRecord.object != undefined) {
+      this.defaultHandler = Object.keys(
+        definitionRecord.object as { [s: string]: unknown }
+      )[0];
+    } else if (definitionRecord && definitionRecord.workflow != undefined) {
+      this.defaultHandler = Object.keys(
+        definitionRecord.workflow as { [s: string]: unknown }
+      )[0];
+    } else {
+      throw new TypeError(
+        "supports only a service or a virtual object or a workflow definition"
+      );
+    }
+  }
 
-    if (entries.length < 2) {
+  public async run(options: RunOptions): Promise<Message[]> {
+    const restateServer = new TestRestateServer();
+    restateServer.bind(this.definition);
+
+    // Sanity check on input messages
+    if (options.input.length < 2) {
       throw new Error(
         "Less than two runtime messages supplied for test. Need to have at least start message and input message."
       );
     }
-
-    if (entries[0].messageType !== START_MESSAGE_TYPE) {
+    if (options.input[0].messageType !== START_MESSAGE_TYPE) {
       throw new Error("First message has to be start message.");
     }
 
     // Get the index of where the completion messages start in the entries list
-    const firstCompletionIndex = entries.findIndex(
+    const firstCompletionIndex = options.input.findIndex(
       (value) =>
         value.messageType === COMPLETION_MESSAGE_TYPE ||
         value.messageType === ENTRY_ACK_MESSAGE_TYPE
     );
 
     // The last message of the replay is the one right before the first completion
     const endOfReplay =
-      firstCompletionIndex !== -1 ? firstCompletionIndex : entries.length;
+      firstCompletionIndex !== -1 ? firstCompletionIndex : options.input.length;
 
-    const msg = entries[0];
-    // We need to set the right number for known entries. Copy the rest
-    const startEntry = msg.message as StartMessage;
-    entries[0] = new Message(
-      msg.messageType,
+    // --- Patch StartMessage with the right number of entries
+    const startMsg = options.input[0];
+    const startEntry = startMsg.message as StartMessage;
+    options.input[0] = new Message(
+      startMsg.messageType,
       new StartMessage({
         id: startEntry.id,
         debugId: startEntry.debugId,
@@ -102,13 +166,18 @@ export class TestDriver implements Connection {
         partialState: startEntry.partialState,
         key: startEntry.key,
       }),
-      msg.completed,
-      msg.requiresAck
+      startMsg.completed,
+      startMsg.requiresAck
     );
 
-    const replayMessages = entries.slice(0, endOfReplay);
-    this.completionMessages = entries.slice(endOfReplay);
-
+    // TODO the production code here is doing some bad assumption,
+    //  by assuming that during the initial replay phase no CompletionMessages are sent.
+    //  Although this is currently correct, it is correct only due to how the runtime is implemented,
+    //  and might not be generally true if we change the runtime.
+    //  This should probably be fixed in the production code, and subsequently the test should
+    //  stop splitting the input messages here.
+    const replayMessages = options.input.slice(0, endOfReplay);
+    const completionMessages = options.input.slice(endOfReplay);
     if (
       replayMessages.filter(
         (value) =>
@@ -120,9 +189,8 @@ export class TestDriver implements Connection {
         "You cannot interleave replay messages with completion or ack messages. First define the replay messages, then the completion messages."
       );
     }
-
     if (
-      this.completionMessages.filter(
+      completionMessages.filter(
         (value) =>
           value.messageType !== COMPLETION_MESSAGE_TYPE &&
           value.messageType !== ENTRY_ACK_MESSAGE_TYPE
@@ -133,90 +201,76 @@ export class TestDriver implements Connection {
       );
     }
 
-    const method = this.restateServer
-      .componentByName("greeter")
+    const method = restateServer
+      .componentByName(options.service ? options.service : this.defaultService)
       ?.handlerMatching({
-        componentName: "greeter",
-        handlerName: "greet",
+        componentName: options.service ? options.service : this.defaultService,
+        handlerName: options.handler ? options.handler : this.defaultHandler,
       });
-
     if (!method) {
-      throw new Error("Something is wrong with the test setup");
+      throw new Error("Can't find the handler to execute");
     }
 
     const invocationBuilder = new InvocationBuilder(method);
     replayMessages.forEach((el) => invocationBuilder.handleMessage(el));
     const invocation = invocationBuilder.build();
 
-    this.stateMachine = new StateMachine(
-      this,
+    const testConnection = new TestConnection();
+    const stateMachine = new StateMachine(
+      testConnection,
       invocation,
-      HandlerKind.EXCLUSIVE,
+      method.kind(),
       // eslint-disable-next-line @typescript-eslint/no-unsafe-argument
-      (
-        this.restateServer as unknown as { builder: EndpointBuilder }
-      ).builder.logger,
+      (restateServer as unknown as { builder: EndpointBuilder }).builder.logger,
       invocation.inferLoggerContext()
     );
-  }
 
-  headers(): ReadonlyMap<string, string | string[] | undefined> {
-    return new Map();
-  }
-
-  async run(): Promise<Message[]> {
-    const completed = this.stateMachine.invoke();
+    const completed = stateMachine.invoke();
 
     // we send the completions here. Because we don't await the messages that we send the completions for,
     // we enqueue those completions in the event loop, so they get processed when everything else is done.
     // This is highly fragile!!!
-    this.completionMessages.forEach((el) => {
-      setTimeout(() => this.stateMachine.handleMessage(el));
+    completionMessages.forEach((el) => {
+      setTimeout(() => stateMachine.handleMessage(el));
     });
     // Set the input channel to closed a bit after sending the messages
     // to make the service finish up the work it can do and suspend or send back a response.
-    setTimeout(() => this.stateMachine.handleInputClosed());
+    setTimeout(() => stateMachine.handleInputClosed());
 
     await completed;
 
-    return Promise.resolve(this.result);
+    return Promise.resolve(testConnection.sentMessages());
+  }
+}
+
+// eslint-disable-next-line @typescript-eslint/no-redundant-type-constituents
+export function testService<N extends string, T>(
+  definition:
+    | ServiceDefinition<N, T>
+    | VirtualObjectDefinition<N, T>
+    | WorkflowDefinition<N, T>
+): UUT<N, T> {
+  return new UUT<N, T>(definition);
+}
+
+class TestConnection implements Connection {
+  private result: Message[] = [];
+
+  headers(): ReadonlyMap<string, string | string[] | undefined> {
+    return new Map();
   }
 
   send(msg: Message): Promise<void> {
     this.result.push(msg);
-    (
-      this.restateServer as unknown as { builder: EndpointBuilder }
-    ).builder.rlog.debug(
-      `Adding result to the result array. Message type: ${
-        msg.messageType
-      }, message:
-        ${
-          msg.message instanceof Uint8Array
-            ? (msg.message as Uint8Array).toString()
-            : formatMessageAsJson(msg.message)
-        }`
-    );
     return Promise.resolve();
   }
 
-  onClose() {
-    // nothing to do
-  }
-
   async end(): Promise<void> {
     // nothing to do
     return Promise.resolve();
   }
 
-  onError() {
-    // nothing to do
+  sentMessages(): Message[] {
+    return this.result;
   }
 }
-
-/**
- * This class' only purpose is to make certain methods accessible in tests.
- * Those methods are otherwise protected, to reduce the public interface and
- * make it simpler for users to understand what methods are relevant for them,
- * and which ones are not.
- */
-class TestRestateServer extends NodeEndpoint {}
