From 36357a277a730b86c98f0f60717e0c09b7821740 Mon Sep 17 00:00:00 2001
From: Jack Kleeman <jackkleeman@gmail.com>
Date: Mon, 4 Nov 2024 11:21:44 +0000
Subject: [PATCH] Simplify container factory

---
 packages/restate-sdk-examples/package.json    |  3 +-
 .../test/testcontainers.test.ts               |  8 ++--
 .../src/public_api.ts                         |  5 +--
 .../src/restate_test_environment.ts           | 37 +++++++------------
 4 files changed, 20 insertions(+), 33 deletions(-)

diff --git a/packages/restate-sdk-examples/package.json b/packages/restate-sdk-examples/package.json
index cf347b4..ff9c1dd 100644
--- a/packages/restate-sdk-examples/package.json
+++ b/packages/restate-sdk-examples/package.json
@@ -40,7 +40,8 @@
   },
   "devDependencies": {
     "tsx": "^4.15.7",
-    "@restatedev/restate-sdk-testcontainers": "^1.4.0"
+    "@restatedev/restate-sdk-testcontainers": "^1.4.0",
+    "testcontainers": "^10.4.0"
   },
   "engines": {
     "node": ">= 18.13"
diff --git a/packages/restate-sdk-examples/test/testcontainers.test.ts b/packages/restate-sdk-examples/test/testcontainers.test.ts
index 122f972..fba8fc2 100644
--- a/packages/restate-sdk-examples/test/testcontainers.test.ts
+++ b/packages/restate-sdk-examples/test/testcontainers.test.ts
@@ -9,14 +9,12 @@
  * https://github.com/restatedev/sdk-typescript/blob/main/LICENSE
  */
 
-import {
-  RestateTestEnvironment,
-  restateContainerFactory,
-} from "@restatedev/restate-sdk-testcontainers";
+import { RestateTestEnvironment } from "@restatedev/restate-sdk-testcontainers";
 import { counter } from "../src/object.js";
 import * as clients from "@restatedev/restate-sdk-clients";
 import * as core from "@restatedev/restate-sdk-core";
 import { describe, it, beforeAll, afterAll, expect } from "vitest";
+import { GenericContainer } from "testcontainers";
 
 describe("ExampleObject", () => {
   let restateTestEnvironment: RestateTestEnvironment;
@@ -129,7 +127,7 @@ describe("Custom testcontainer config", () => {
     restateTestEnvironment = await RestateTestEnvironment.start(
       (restateServer) => restateServer.bind(counter),
       () =>
-        restateContainerFactory("restatedev/restate:1.1")()
+        new GenericContainer("restatedev/restate:1.1")
           .withEnvironment({ RESTATE_LOG_FORMAT: "json" })
           .withLogConsumer((stream) => {
             // eslint-disable-next-line no-console
diff --git a/packages/restate-sdk-testcontainers/src/public_api.ts b/packages/restate-sdk-testcontainers/src/public_api.ts
index 027e4cf..95fd038 100644
--- a/packages/restate-sdk-testcontainers/src/public_api.ts
+++ b/packages/restate-sdk-testcontainers/src/public_api.ts
@@ -9,7 +9,4 @@
  * https://github.com/restatedev/sdk-typescript/blob/main/LICENSE
  */
 
-export {
-  RestateTestEnvironment,
-  restateContainerFactory,
-} from "./restate_test_environment.js";
+export { RestateTestEnvironment } from "./restate_test_environment.js";
diff --git a/packages/restate-sdk-testcontainers/src/restate_test_environment.ts b/packages/restate-sdk-testcontainers/src/restate_test_environment.ts
index 3b23ff0..d12b9a8 100644
--- a/packages/restate-sdk-testcontainers/src/restate_test_environment.ts
+++ b/packages/restate-sdk-testcontainers/src/restate_test_environment.ts
@@ -50,9 +50,18 @@ async function prepareRestateEndpoint(
 // Prepare the restate testcontainer
 async function prepareRestateTestContainer(
   restateServerPort: number,
-  restateContainerFactory: () => GenericContainer = () => new GenericContainer("docker.io/restatedev/restate:1.1")
+  restateContainerFactory: () => GenericContainer
 ): Promise<StartedTestContainer> {
-  const restateContainer = restateContainerFactory();
+  const restateContainer = restateContainerFactory()
+    // Expose ports
+    .withExposedPorts(8080, 9070)
+    // Wait start on health checks
+    .withWaitStrategy(
+      Wait.forAll([
+        Wait.forHttp("/restate/health", 8080),
+        Wait.forHttp("/health", 9070),
+      ])
+    );
 
   // This MUST be executed before starting the restate container
   // Expose host port to access the restate server
@@ -102,23 +111,6 @@ async function prepareRestateTestContainer(
   }
 }
 
-export function restateContainerFactory(image: string): () => GenericContainer {
-  return () => {
-    const restateContainer = new GenericContainer(image)
-      // Expose ports
-      .withExposedPorts(8080, 9070)
-      // Wait start on health checks
-      .withWaitStrategy(
-        Wait.forAll([
-          Wait.forHttp("/restate/health", 8080),
-          Wait.forHttp("/health", 9070),
-        ])
-      );
-
-    return restateContainer;
-  };
-}
-
 export class RestateTestEnvironment {
   constructor(
     readonly startedRestateHttpServer: http2.Http2Server,
@@ -154,16 +146,15 @@ export class RestateTestEnvironment {
 
   public static async start(
     mountServicesFn: (server: restate.RestateEndpoint) => void,
-    _restateContainerFactory: () => GenericContainer = restateContainerFactory(
-      "restatedev/restate:1.1"
-    )
+    restateContainerFactory: () => GenericContainer = () =>
+      new GenericContainer("restatedev/restate:1.1")
   ): Promise<RestateTestEnvironment> {
     const startedRestateHttpServer = await prepareRestateEndpoint(
       mountServicesFn
     );
     const startedRestateContainer = await prepareRestateTestContainer(
       (startedRestateHttpServer.address() as net.AddressInfo).port,
-      _restateContainerFactory
+      restateContainerFactory
     );
     return new RestateTestEnvironment(
       startedRestateHttpServer,