Reassign runningWorkflows workflow ID in test environment

internal/internal_workflow_testsuite.go

@@ -2431,6 +2431,10 @@ func (env *testWorkflowEnvironmentImpl) setStartWorkflowOptions(options StartWor
 		wf.WorkflowTaskTimeout = options.WorkflowTaskTimeout
 	}
 	if len(options.ID) > 0 {
+		// Reassign the ID in running Workflows so SignalWorkflowByID can find the workflow
+		originalID := wf.WorkflowExecution.ID
+		env.runningWorkflows[options.ID] = env.runningWorkflows[wf.WorkflowExecution.ID]
+		delete(env.runningWorkflows, originalID)
 		wf.WorkflowExecution.ID = options.ID
 	}
 	if len(options.TaskQueue) > 0 {

internal/workflow_testsuite_test.go

@@ -27,11 +27,12 @@ package internal
 import (
 	"context"
 	"errors"
-	"github.com/stretchr/testify/assert"
 	"strings"
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
+
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 )
@@ -227,6 +228,27 @@ func TestWorkflowIDInsideTestWorkflow(t *testing.T) {
 	require.Equal(t, "id is: my-workflow-id", str)
 }
 
+func TestWorkflowIDSignalWorkflowByID(t *testing.T) {
+	var suite WorkflowTestSuite
+	// Test SignalWorkflowByID works with custom ID
+	env := suite.NewTestWorkflowEnvironment()
+	env.RegisterDelayedCallback(func() {
+		err := env.SignalWorkflowByID("my-workflow-id", "signal", "payload")
+		require.NoError(t, err)
+	}, time.Second)
+
+	env.SetStartWorkflowOptions(StartWorkflowOptions{ID: "my-workflow-id"})
+	env.ExecuteWorkflow(func(ctx Context) (string, error) {
+		var result string
+		GetSignalChannel(ctx, "signal").Receive(ctx, &result)
+		return "id is: " + GetWorkflowInfo(ctx).WorkflowExecution.ID, nil
+	})
+	require.NoError(t, env.GetWorkflowError())
+	var str string
+	require.NoError(t, env.GetWorkflowResult(&str))
+	require.Equal(t, "id is: my-workflow-id", str)
+}
+
 func TestWorkflowStartTimeInsideTestWorkflow(t *testing.T) {
 	var suite WorkflowTestSuite
 	env := suite.NewTestWorkflowEnvironment()