From 7df44bfa0df8853e802b6e172bbb63bd40658942 Mon Sep 17 00:00:00 2001
From: Quinn Klassen <klassenq@gmail.com>
Date: Sat, 15 Jul 2023 15:07:43 -0700
Subject: [PATCH] Reassign runningWorkflows workflow ID in test environment

---
 internal/internal_workflow_testsuite.go |  4 ++++
 internal/workflow_testsuite_test.go     | 24 +++++++++++++++++++++++-
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/internal/internal_workflow_testsuite.go b/internal/internal_workflow_testsuite.go
index 4d7b368ec..56102597b 100644
--- a/internal/internal_workflow_testsuite.go
+++ b/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 {
diff --git a/internal/workflow_testsuite_test.go b/internal/workflow_testsuite_test.go
index 4e91c77fb..6068e51d2 100644
--- a/internal/workflow_testsuite_test.go
+++ b/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()