feat: Ability to specify initial Key Value Entity state in testkit

akka-javasdk-testkit/src/main/java/akka/javasdk/testkit/KeyValueEntityTestKit.java

@@ -30,6 +30,8 @@
  */
 public class KeyValueEntityTestKit<S, E extends KeyValueEntity<S>> {
 
+  public static final String DEFAULT_TEST_ENTITY_ID = "testkit-entity-id";
+
   private final Class<?> stateClass;
   private S state;
   private boolean deleted;
@@ -46,6 +48,16 @@ private KeyValueEntityTestKit(String entityId, E entity) {
     this.deleted = false;
   }
 
+  private KeyValueEntityTestKit(String entityId, E entity, S initialState) {
+    this.entityId = entityId;
+    this.entity = entity;
+    this.stateClass = Reflect.keyValueEntityStateType(entity.getClass());
+    this.emptyState = entity.emptyState();
+    verifySerDerWithExpectedType(stateClass, initialState, entity);
+    this.state = initialState;
+    this.deleted = false;
+  }
+
   /**
    * Creates a new testkit instance from a KeyValueEntity Supplier.
    *
@@ -63,7 +75,7 @@ public static <S, E extends KeyValueEntity<S>> KeyValueEntityTestKit<S, E> of(
    */
   public static <S, E extends KeyValueEntity<S>> KeyValueEntityTestKit<S, E> of(
       Function<KeyValueEntityContext, E> entityFactory) {
-    return of("testkit-entity-id", entityFactory);
+    return of(DEFAULT_TEST_ENTITY_ID, entityFactory);
   }
 
   /** Creates a new testkit instance from a user defined entity id and a KeyValueEntity Supplier. */
@@ -82,6 +94,52 @@ public static <S, E extends KeyValueEntity<S>> KeyValueEntityTestKit<S, E> of(
     return new KeyValueEntityTestKit<>(entityId, entityFactory.apply(context));
   }
 
+  /**
+   * Creates a new testkit instance from a KeyValueEntity Supplier, starting with the given initial
+   * state instead of the entity's {@code emptyState()}. Subsequent {@code deleteEntity()} effects
+   * still reset the state to the entity's {@code emptyState()}.
+   *
+   * <p>A default test entity id will be automatically provided.
+   */
+  public static <S, E extends KeyValueEntity<S>> KeyValueEntityTestKit<S, E> ofEntityWithState(
+      Supplier<E> entityFactory, S initialState) {
+    return ofEntityWithState(ctx -> entityFactory.get(), initialState);
+  }
+
+  /**
+   * Creates a new testkit instance from a function KeyValueEntityContext to KeyValueEntity,
+   * starting with the given initial state instead of the entity's {@code emptyState()}. Subsequent
+   * {@code deleteEntity()} effects still reset the state to the entity's {@code emptyState()}.
+   *
+   * <p>A default test entity id will be automatically provided.
+   */
+  public static <S, E extends KeyValueEntity<S>> KeyValueEntityTestKit<S, E> ofEntityWithState(
+      Function<KeyValueEntityContext, E> entityFactory, S initialState) {
+    return ofEntityWithState(DEFAULT_TEST_ENTITY_ID, entityFactory, initialState);
+  }
+
+  /**
+   * Creates a new testkit instance from a user defined entity id and a KeyValueEntity Supplier,
+   * starting with the given initial state instead of the entity's {@code emptyState()}. Subsequent
+   * {@code deleteEntity()} effects still reset the state to the entity's {@code emptyState()}.
+   */
+  public static <S, E extends KeyValueEntity<S>> KeyValueEntityTestKit<S, E> ofEntityWithState(
+      String entityId, Supplier<E> entityFactory, S initialState) {
+    return ofEntityWithState(entityId, ctx -> entityFactory.get(), initialState);
+  }
+
+  /**
+   * Creates a new testkit instance from a user defined entity id and a function
+   * KeyValueEntityContext to KeyValueEntity, starting with the given initial state instead of the
+   * entity's {@code emptyState()}. Subsequent {@code deleteEntity()} effects still reset the state
+   * to the entity's {@code emptyState()}.
+   */
+  public static <S, E extends KeyValueEntity<S>> KeyValueEntityTestKit<S, E> ofEntityWithState(
+      String entityId, Function<KeyValueEntityContext, E> entityFactory, S initialState) {
+    TestKitKeyValueEntityContext context = new TestKitKeyValueEntityContext(entityId);
+    return new KeyValueEntityTestKit<>(entityId, entityFactory.apply(context), initialState);
+  }
+
   /**
    * @return The current state of the key value entity under test
    */

akka-javasdk-testkit/src/test/java/akka/javasdk/testkit/keyvalueentity/CounterKeyValueEntityTest.java

@@ -53,6 +53,25 @@ public void testIncreaseWithNegativeValue() {
     assertEquals(result.getError(), "Can't increase with a negative value");
   }
 
+  @Test
+  public void testInitialState() {
+    KeyValueEntityTestKit<Integer, CounterValueEntity> testKit =
+        KeyValueEntityTestKit.ofEntityWithState(ctx -> new CounterValueEntity(), 42);
+    assertEquals(42, testKit.getState());
+    KeyValueEntityResult<String> result = testKit.method(CounterValueEntity::increaseBy).invoke(8);
+    assertTrue(result.isReply());
+    assertEquals(50, testKit.getState());
+  }
+
+  @Test
+  public void testInitialStateResetOnDelete() {
+    KeyValueEntityTestKit<Integer, CounterValueEntity> testKit =
+        KeyValueEntityTestKit.ofEntityWithState(ctx -> new CounterValueEntity(), 42);
+    testKit.method(CounterValueEntity::delete).invoke();
+    assertTrue(testKit.isDeleted());
+    assertEquals(0, testKit.getState());
+  }
+
   @Test
   public void testDeleteValueEntity() {
     KeyValueEntityTestKit<Integer, CounterValueEntity> testKit =