From a6cdcd74afcaba5294b5e97132cc82c1cb8f490f Mon Sep 17 00:00:00 2001
From: Quinn Klassen <klassenq@gmail.com>
Date: Thu, 12 Oct 2023 13:55:38 -0700
Subject: [PATCH] Add UpdateWorkflowByID to TestWorkflowEnvironment (#1264)

Add UpdateWorkflowByID to TestWorkflowEnvironment
---
 internal/internal_workflow_testsuite.go | 16 ++++++++++++
 internal/workflow_testsuite.go          |  9 +++++--
 internal/workflow_testsuite_test.go     | 34 +++++++++++++++++++++++++
 3 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/internal/internal_workflow_testsuite.go b/internal/internal_workflow_testsuite.go
index c1be8b4d7..740b65a68 100644
--- a/internal/internal_workflow_testsuite.go
+++ b/internal/internal_workflow_testsuite.go
@@ -2382,6 +2382,22 @@ func (env *testWorkflowEnvironmentImpl) updateWorkflow(name string, id string, u
 	env.updateHandler(name, id, data, nil, uc)
 }
 
+func (env *testWorkflowEnvironmentImpl) updateWorkflowByID(workflowID, name, id string, uc UpdateCallbacks, args ...interface{}) error {
+	if workflowHandle, ok := env.runningWorkflows[workflowID]; ok {
+		if workflowHandle.handled {
+			return serviceerror.NewNotFound(fmt.Sprintf("Workflow %v already completed", workflowID))
+		}
+		data, err := encodeArgs(env.GetDataConverter(), args)
+		if err != nil {
+			panic(err)
+		}
+		env.updateHandler(name, id, data, nil, uc)
+		return nil
+	}
+
+	return serviceerror.NewNotFound(fmt.Sprintf("Workflow %v not exists", workflowID))
+}
+
 func (env *testWorkflowEnvironmentImpl) queryWorkflowByID(workflowID, queryType string, args ...interface{}) (converter.EncodedValue, error) {
 	if workflowHandle, ok := env.runningWorkflows[workflowID]; ok {
 		data, err := encodeArgs(workflowHandle.env.GetDataConverter(), args)
diff --git a/internal/workflow_testsuite.go b/internal/workflow_testsuite.go
index aef295aa0..3dc57f1a3 100644
--- a/internal/workflow_testsuite.go
+++ b/internal/workflow_testsuite.go
@@ -853,8 +853,13 @@ func (e *TestWorkflowEnvironment) QueryWorkflow(queryType string, args ...interf
 }
 
 // UpdateWorkflow sends an update to the currently running workflow.
-func (e *TestWorkflowEnvironment) UpdateWorkflow(name string, id string, uc UpdateCallbacks, args ...interface{}) {
-	e.impl.updateWorkflow(name, id, uc, args...)
+func (e *TestWorkflowEnvironment) UpdateWorkflow(updateName, updateID string, uc UpdateCallbacks, args ...interface{}) {
+	e.impl.updateWorkflow(updateName, updateID, uc, args...)
+}
+
+// UpdateWorkflowByID sends an update to a running workflow by its ID.
+func (e *TestWorkflowEnvironment) UpdateWorkflowByID(workflowID, updateName, updateID string, uc UpdateCallbacks, args interface{}) error {
+	return e.impl.updateWorkflowByID(workflowID, updateName, updateID, uc, args)
 }
 
 // QueryWorkflowByID queries a child workflow by its ID and returns the result synchronously
diff --git a/internal/workflow_testsuite_test.go b/internal/workflow_testsuite_test.go
index 6068e51d2..4496797d8 100644
--- a/internal/workflow_testsuite_test.go
+++ b/internal/workflow_testsuite_test.go
@@ -249,6 +249,40 @@ func TestWorkflowIDSignalWorkflowByID(t *testing.T) {
 	require.Equal(t, "id is: my-workflow-id", str)
 }
 
+func TestWorkflowIDUpdateWorkflowByID(t *testing.T) {
+	var suite WorkflowTestSuite
+	// Test UpdateWorkflowByID works with custom ID
+	env := suite.NewTestWorkflowEnvironment()
+	env.RegisterDelayedCallback(func() {
+		err := env.UpdateWorkflowByID("my-workflow-id", "update", "id", &updateCallback{
+			reject: func(err error) {
+				require.Fail(t, "update should not be rejected")
+			},
+			accept:   func() {},
+			complete: func(interface{}, error) {},
+		}, "input")
+		require.NoError(t, err)
+	}, time.Second)
+
+	env.SetStartWorkflowOptions(StartWorkflowOptions{ID: "my-workflow-id"})
+	env.ExecuteWorkflow(func(ctx Context) (string, error) {
+		var result string
+		err := SetUpdateHandler(ctx, "update", func(ctx Context, input string) error {
+			result = input
+			return nil
+		}, UpdateHandlerOptions{})
+		if err != nil {
+			return "", err
+		}
+		err = Await(ctx, func() bool { return result != "" })
+		return result, err
+	})
+	require.NoError(t, env.GetWorkflowError())
+	var str string
+	require.NoError(t, env.GetWorkflowResult(&str))
+	require.Equal(t, "input", str)
+}
+
 func TestWorkflowStartTimeInsideTestWorkflow(t *testing.T) {
 	var suite WorkflowTestSuite
 	env := suite.NewTestWorkflowEnvironment()