Add getState and setState methods in RestateTestEnvironment (#207)

templates/typescript-testing/package.json

@@ -20,6 +20,7 @@
   "devDependencies": {
     "@types/jest": "^29.4.0",
     "@types/node": "^20.14.2",
+    "apache-arrow": "^18.0.0",
     "esbuild": "^0.21.5",
     "testcontainers": "^10.4.0",
     "ts-jest": "^29.0.5",

templates/typescript-testing/src/app.ts

@@ -1,5 +1,6 @@
 import * as restate from "@restatedev/restate-sdk";
 import {exampleService} from "./example_service";
+import {exampleObject} from "./example_object";
 
 // Template of a Restate service and handler
 //
@@ -10,4 +11,5 @@ import {exampleService} from "./example_service";
 restate
   .endpoint()
   .bind( exampleService )
+  .bind( exampleObject )
   .listen(9080);

templates/typescript-testing/src/example_object.ts

@@ -0,0 +1,17 @@
+import * as restate from "@restatedev/restate-sdk";
+
+// Template of a Restate virtual object and handler
+//
+// Have a look at the TS QuickStart to learn how to run this: https://docs.restate.dev/get_started/quickstart?sdk=ts
+//
+
+export const exampleObject = restate.object({
+    name: "ExampleObject",
+    handlers: {
+        greet: async (ctx: restate.ObjectContext) => {
+          const count = (await ctx.get<number>("count")) ?? 0;
+          ctx.set("count", count + 1);
+          return `Hello ${ctx.key}! Counter: ${count}`;
+        }
+    },
+})

templates/typescript-testing/test/restate_test_environment.ts

@@ -6,6 +6,7 @@ import {
     TestContainers,
     Wait,
 } from "testcontainers";
+import { tableFromIPC } from "apache-arrow";
 import * as http2 from "http2";
 import * as net from "net";
 
@@ -109,6 +110,68 @@ export class RestateTestEnvironment {
         )}`;
     }
 
+    public async setState(
+      service: restate.VirtualObjectDefinition<any, any> | restate.WorkflowDefinition<any, any>,
+      key: string, 
+      newState: {[key: string]: any}) {
+      const res = await fetch(
+          `${this.adminAPIBaseUrl()}/services/${service.name}/state`,
+          {
+              method: "POST",
+              headers: {
+                  "Content-Type": "application/json",
+              },
+              body: JSON.stringify({
+                  object_key: key,
+                  // the endpoint expects a map of key -> bytes as JSON array of numbers
+                  new_state: Object.fromEntries(Object.entries(newState).map(([key, value]) => {
+                    const valueJSON = new TextEncoder().encode(JSON.stringify(value))
+
+                    return [key, Array.from(valueJSON)]
+                  })),
+              }),
+          }
+      );
+
+      if (!res.ok) {
+          const badResponse = await res.text();
+          throw new Error(
+              `Error ${res.status} during modify state: ${badResponse}`
+          );
+      }
+    }
+
+    public async getState(
+      service: restate.VirtualObjectDefinition<any, any> | restate.WorkflowDefinition<any, any>,
+      key: string
+    ): Promise<{[key: string]: any}> {
+      const res = await fetch(
+          `${this.adminAPIBaseUrl()}/query`,
+          {
+              method: "POST",
+              headers: {
+                  "Content-Type": "application/json",
+              },
+              body: JSON.stringify({
+                  query: `SELECT key, value from state where service_name = '${service.name}' and service_key = '${key}';`,
+              }),
+          }
+      );
+
+      if (!res.ok) {
+          const badResponse = await res.text();
+          throw new Error(
+              `Error ${res.status} during read state: ${badResponse}`
+          );
+      }
+
+      const table = (await tableFromIPC(res)).toArray() as { key: string, value: Uint8Array }[];
+
+      return Object.fromEntries(table.map(({key, value}) => {
+        return [key, JSON.parse(new TextDecoder().decode(value))]
+      }))
+    }
+
     public async stop() {
         await this.startedRestateContainer.stop();
         this.startedRestateHttpServer.close();
@@ -126,4 +189,4 @@ export class RestateTestEnvironment {
             startedRestateContainer
         );
     }
-}
+}

templates/typescript-testing/test/test.ts

@@ -1,5 +1,6 @@
 import { RestateTestEnvironment } from "./restate_test_environment";
 import { exampleService } from "../src/example_service";
+import { exampleObject } from "../src/example_object";
 import * as clients from "@restatedev/restate-sdk-clients";
 
 describe("ExampleService", () => {
@@ -28,4 +29,36 @@ describe("ExampleService", () => {
         // Assert the result
         expect(greet).toBe("Hello Sarah!");
     });
-});
+});
+
+describe("ExampleObject", () => {
+    let restateTestEnvironment: RestateTestEnvironment;
+
+    // Deploy Restate and the Service endpoint once for all the tests in this suite
+    beforeAll(async () => {
+        restateTestEnvironment = await RestateTestEnvironment.start(
+            (restateServer) =>
+                restateServer.bind(exampleObject)
+        );
+    }, 20_000);
+
+    // Stop Restate and the Service endpoint
+    afterAll(async () => {
+        if (restateTestEnvironment !== undefined) {
+            await restateTestEnvironment.stop();
+        }
+    });
+
+    it("works", async () => {
+        const rs = clients.connect({url: restateTestEnvironment.baseUrl()});
+        expect(await restateTestEnvironment.getState(exampleObject, "Sarah")).toStrictEqual({})
+
+        await restateTestEnvironment.setState(exampleObject, "Sarah", {count: 123})
+        const greet = await rs.objectClient(exampleObject, "Sarah")
+            .greet();
+
+        // Assert the result
+        expect(greet).toBe("Hello Sarah! Counter: 123");
+        expect(await restateTestEnvironment.getState(exampleObject, "Sarah")).toStrictEqual({count: 124})
+    });
+});