add step idempotency tests

Author: maxdml

dbos/queues_test.go

@@ -634,6 +634,8 @@ func TestWorkerConcurrencyXRecovery(t *testing.T) {
 	// Start the first workflow and wait for it to start
 	workerConcurrencyRecoveryStartEvent1.Wait()
 	workerConcurrencyRecoveryStartEvent1.Clear()
+	// Wait for a few seconds to let the queue runner loop
+	time.Sleep(2 * time.Second)
 
 	// Ensure the 2nd workflow is still ENQUEUED
 	status2, err := handle2.GetStatus()

dbos/workflows_test.go

@@ -335,7 +335,16 @@ func stepRetryAlwaysFailsStep(ctx context.Context, input string) (string, error)
 	return "", fmt.Errorf("always fails - attempt %d", stepRetryAttemptCount)
 }
 
+var stepIdempotencyCounter int
+
+func stepIdempotencyTest(ctx context.Context, input string) (string, error) {
+	stepIdempotencyCounter++
+	fmt.Println("Executing idempotency step:", stepIdempotencyCounter)
+	return input, nil
+}
+
 func stepRetryWorkflow(ctx context.Context, input string) (string, error) {
+	RunAsStep(ctx, stepIdempotencyTest, input)
 	return RunAsStep(ctx, stepRetryAlwaysFailsStep, input,
 		WithStepMaxRetries(5),
 		WithBackoffFactor(2.0),
@@ -400,8 +409,9 @@ func TestSteps(t *testing.T) {
 	})
 
 	t.Run("StepRetryWithExponentialBackoff", func(t *testing.T) {
-		// Reset the global counter before test
+		// Reset the global counters before test
 		stepRetryAttemptCount = 0
+		stepIdempotencyCounter = 0
 
 		// Execute the workflow
 		handle, err := stepRetryWf(context.Background(), "test")
@@ -449,19 +459,24 @@ func TestSteps(t *testing.T) {
 			t.Fatal("failed to get workflow steps:", err)
 		}
 
-		if len(steps) != 1 {
-			t.Fatalf("expected 1 recorded step, got %d", len(steps))
+		if len(steps) != 2 {
+			t.Fatalf("expected 2 recorded step, got %d", len(steps))
 		}
 
-		// Verify the recorded step has the error
-		step := steps[0]
+		// Verify the second step has the error
+		step := steps[1]
 		if step.Error == nil {
 			t.Fatal("expected error in recorded step, got none")
 		}
 
 		if step.Error.Error() != dbosErr.Error() {
 			t.Fatalf("expected recorded step error to match joined error, got '%s', expected '%s'", step.Error.Error(), dbosErr.Error())
 		}
+
+		// Verify the idempotency step was executed only once
+		if stepIdempotencyCounter != 1 {
+			t.Fatalf("expected idempotency step to be executed only once, got %d", stepIdempotencyCounter)
+		}
 	})
 }
 
@@ -980,24 +995,22 @@ func TestScheduledWorkflows(t *testing.T) {
 }
 
 var (
-	sendWf          = WithWorkflow(sendWorkflow)
-	singleSendWf    = WithWorkflow(singleSendWorkflow)
-	receiveWf       = WithWorkflow(receiveWorkflow)
-	receiveWfSteps  = WithWorkflow(receiveWorkflowWithSteps)
-	sendStructWf    = WithWorkflow(sendStructWorkflow)
-	receiveStructWf = WithWorkflow(receiveStructWorkflow)
+	sendWf                       = WithWorkflow(sendWorkflow)
+	receiveWf                    = WithWorkflow(receiveWorkflow)
+	sendStructWf                 = WithWorkflow(sendStructWorkflow)
+	receiveStructWf              = WithWorkflow(receiveStructWorkflow)
+	sendIdempotencyWf            = WithWorkflow(sendIdempotencyWorkflow)
+	sendIdempotencyEvent         = NewEvent()
+	recvIdempotencyWf            = WithWorkflow(receiveIdempotencyWorkflow)
+	receiveIdempotencyStartEvent = NewEvent()
+	receiveIdempotencyStopEvent  = NewEvent()
 )
 
 type sendWorkflowInput struct {
 	DestinationID string
 	Topic         string
 }
 
-var (
-	testRecvStartEvent    = NewEvent()
-	testSendCompleteEvent = NewEvent()
-)
-
 func sendWorkflow(ctx context.Context, input sendWorkflowInput) (string, error) {
 	fmt.Println("Starting send workflow with input:", input)
 	err := Send(ctx, WorkflowSendInput{DestinationID: input.DestinationID, Topic: input.Topic, Message: "message1"})
@@ -1016,20 +1029,6 @@ func sendWorkflow(ctx context.Context, input sendWorkflowInput) (string, error)
 	return "", nil
 }
 
-func singleSendWorkflow(ctx context.Context, input sendWorkflowInput) (string, error) {
-	fmt.Println("Starting send workflow with input:", input)
-
-	testSendCompleteEvent.Set()
-	err := Send(ctx, WorkflowSendInput{DestinationID: input.DestinationID, Topic: input.Topic, Message: "message1"})
-	if err != nil {
-		fmt.Printf("Send failed with error: %v\n", err)
-		return "", err
-	}
-
-	fmt.Println("Sending message on topic:", input.Topic, "to destination:", input.DestinationID)
-	return "", nil
-}
-
 func receiveWorkflow(ctx context.Context, topic string) (string, error) {
 	msg1, err := Recv[string](ctx, WorkflowRecvInput{Topic: topic, Timeout: 3 * time.Second})
 	if err != nil {
@@ -1046,25 +1045,6 @@ func receiveWorkflow(ctx context.Context, topic string) (string, error) {
 	return msg1 + "-" + msg2 + "-" + msg3, nil
 }
 
-func receiveWorkflowWithSteps(ctx context.Context, input string) (string, error) {
-	_, err := RunAsStep(ctx, simpleStep, input)
-	if err != nil {
-		return "", err
-	}
-
-	_, err = RunAsStep(ctx, simpleStep, input)
-	if err != nil {
-		return "", err
-	}
-
-	testRecvStartEvent.Set()
-	msg1, err := Recv[string](ctx, WorkflowRecvInput{Topic: input, Timeout: 1 * time.Second})
-	if err != nil {
-		return "", err
-	}
-	return msg1, nil
-}
-
 func sendStructWorkflow(ctx context.Context, input sendWorkflowInput) (string, error) {
 	testStruct := sendRecvType{Value: "test-struct-value"}
 	err := Send(ctx, WorkflowSendInput{DestinationID: input.DestinationID, Topic: input.Topic, Message: testStruct})
@@ -1075,6 +1055,25 @@ func receiveStructWorkflow(ctx context.Context, topic string) (sendRecvType, err
 	return Recv[sendRecvType](ctx, WorkflowRecvInput{Topic: topic, Timeout: 3 * time.Second})
 }
 
+func sendIdempotencyWorkflow(ctx context.Context, input sendWorkflowInput) (string, error) {
+	err := Send(ctx, WorkflowSendInput{DestinationID: input.DestinationID, Topic: input.Topic, Message: "m1"})
+	if err != nil {
+		return "", err
+	}
+	sendIdempotencyEvent.Wait()
+	return "idempotent-send-completed", nil
+}
+
+func receiveIdempotencyWorkflow(ctx context.Context, topic string) (string, error) {
+	msg, err := Recv[string](ctx, WorkflowRecvInput{Topic: topic, Timeout: 3 * time.Second})
+	if err != nil {
+		return "", err
+	}
+	receiveIdempotencyStartEvent.Set()
+	receiveIdempotencyStopEvent.Wait()
+	return msg, nil
+}
+
 type sendRecvType struct {
 	Value string
 }
@@ -1158,31 +1157,6 @@ func TestSendRecv(t *testing.T) {
 		}
 	})
 
-	t.Run("SendUsesCorrectWorkflowIDForIdempotency", func(t *testing.T) {
-		receiveHandle, err := receiveWfSteps(context.Background(), "test-topic")
-		if err != nil {
-			t.Fatalf("failed to start receive workflow: %v", err)
-		}
-
-		testRecvStartEvent.Wait()
-
-		sendHandle, err := singleSendWf(context.Background(), sendWorkflowInput{
-			DestinationID: receiveHandle.GetWorkflowID(),
-			Topic:         "test-topic",
-		})
-
-		if err != nil {
-			t.Fatalf("failed to send message: %v", err)
-		}
-
-		testSendCompleteEvent.Wait()
-
-		_, err = sendHandle.GetResult(context.Background())
-		if err != nil {
-			t.Fatalf("send workflow failed: %v", err)
-		}
-	})
-
 	t.Run("SendRecvCustomStruct", func(t *testing.T) {
 		// Start the receive workflow
 		receiveHandle, err := receiveStructWf(context.Background(), "struct-topic")
@@ -1310,5 +1284,64 @@ func TestSendRecv(t *testing.T) {
 			t.Fatalf("expected error message to contain %q, but got %q", expectedMessagePart, err.Error())
 		}
 	})
+	t.Run("SendRecvIdempotency", func(t *testing.T) {
+		// Start the receive workflow and wait for it to be ready
+		receiveHandle, err := recvIdempotencyWf(context.Background(), "idempotency-topic")
+		if err != nil {
+			t.Fatalf("failed to start receive idempotency workflow: %v", err)
+		}
+
+		// Send the message to the receive workflow
+		sendHandle, err := sendIdempotencyWf(context.Background(), sendWorkflowInput{
+			DestinationID: receiveHandle.GetWorkflowID(),
+			Topic:         "idempotency-topic",
+		})
+		if err != nil {
+			t.Fatalf("failed to send idempotency message: %v", err)
+		}
 
+		// Wait for the receive workflow to have received the message
+		receiveIdempotencyStartEvent.Wait()
+
+		// Attempt recovering both workflows. There should be only 2 steps recorded after recovery.
+		recoveredHandles, err := recoverPendingWorkflows(context.Background(), []string{"local"})
+		if err != nil {
+			t.Fatalf("failed to recover pending workflows: %v", err)
+		}
+		if len(recoveredHandles) != 2 {
+			t.Fatalf("expected 2 recovered handles, got %d", len(recoveredHandles))
+		}
+		steps, err := getExecutor().systemDB.GetWorkflowSteps(context.Background(), sendHandle.GetWorkflowID())
+		if err != nil {
+			t.Fatalf("failed to get steps for send idempotency workflow: %v", err)
+		}
+		if len(steps) != 1 {
+			t.Fatalf("expected 1 step in send idempotency workflow, got %d", len(steps))
+		}
+		steps, err = getExecutor().systemDB.GetWorkflowSteps(context.Background(), receiveHandle.GetWorkflowID())
+		if err != nil {
+			t.Fatalf("failed to get steps for receive idempotency workflow: %v", err)
+		}
+		if len(steps) != 1 {
+			t.Fatalf("expected 1 step in receive idempotency workflow, got %d", len(steps))
+		}
+
+		// Unblock the workflows to complete
+		receiveIdempotencyStopEvent.Set()
+		result, err := receiveHandle.GetResult(context.Background())
+		if err != nil {
+			t.Fatalf("failed to get result from receive idempotency workflow: %v", err)
+		}
+		if result != "m1" {
+			t.Fatalf("expected result to be 'm1', got '%s'", result)
+		}
+		sendIdempotencyEvent.Set()
+		result, err = sendHandle.GetResult(context.Background())
+		if err != nil {
+			t.Fatalf("failed to get result from send idempotency workflow: %v", err)
+		}
+		if result != "idempotent-send-completed" {
+			t.Fatalf("expected result to be 'idempotent-send-completed', got '%s'", result)
+		}
+	})
 }