adding set & get data attributes from RPC. Missing Synchronous/lockin… (#586)

Co-authored-by: Katie Atrops <[email protected]>

Author: ktrops

integ/rpc_external_storage_test.go

@@ -0,0 +1,164 @@
+package integ
+
+import (
+	"context"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/indeedeng/iwf/gen/iwfidl"
+	rpcStorage "github.com/indeedeng/iwf/integ/workflow/rpc-external-storage"
+	"github.com/indeedeng/iwf/service"
+	"github.com/indeedeng/iwf/service/common/ptr"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRpcExternalStorageNonLockingTemporal(t *testing.T) {
+	if !*temporalIntegTest {
+		t.Skip()
+	}
+	doTestRpcExternalStorage(t, service.BackendTypeTemporal, false)
+}
+
+func TestRpcExternalStorageSynchronousUpdateTemporal(t *testing.T) {
+	if !*temporalIntegTest {
+		t.Skip()
+	}
+	doTestRpcExternalStorage(t, service.BackendTypeTemporal, true)
+}
+
+func TestRpcExternalStorageNonLockingCadence(t *testing.T) {
+	if !*cadenceIntegTest {
+		t.Skip()
+	}
+	// Cadence doesn't support synchronous updates
+	doTestRpcExternalStorage(t, service.BackendTypeCadence, false)
+}
+
+func doTestRpcExternalStorage(t *testing.T, backendType service.BackendType, useLocking bool) {
+	assertions := assert.New(t)
+
+	// start test workflow server
+	wfHandler := rpcStorage.NewHandler()
+	closeFunc1 := startWorkflowWorkerWithRpc(wfHandler, t)
+	defer closeFunc1()
+
+	// Start IWF service with external storage enabled
+	_, closeFunc2 := startIwfServiceByConfig(IwfServiceTestConfig{
+		BackendType:     backendType,
+		S3TestThreshold: 100, // Set low threshold so large data gets stored in S3
+	})
+	defer closeFunc2()
+
+	// create client
+	apiClient := iwfidl.NewAPIClient(&iwfidl.Configuration{
+		Servers: []iwfidl.ServerConfiguration{
+			{
+				URL: "http://localhost:" + testIwfServerPort,
+			},
+		},
+	})
+
+	// start a workflow
+	wfId := rpcStorage.WorkflowType + strconv.Itoa(int(time.Now().UnixNano()))
+	wfInput := &iwfidl.EncodedObject{
+		Encoding: iwfidl.PtrString("json"),
+		Data:     iwfidl.PtrString("\"start-input\""),
+	}
+	req := apiClient.DefaultApi.ApiV1WorkflowStartPost(context.Background())
+	_, httpResp, err := req.WorkflowStartRequest(iwfidl.WorkflowStartRequest{
+		WorkflowId:             wfId,
+		IwfWorkflowType:        rpcStorage.WorkflowType,
+		WorkflowTimeoutSeconds: 30,
+		IwfWorkerUrl:           "http://localhost:" + testWorkflowServerPort,
+		StartStateId:           ptr.Any(rpcStorage.State1),
+		StateInput:             wfInput,
+	}).Execute()
+	failTestAtHttpError(err, httpResp, t)
+
+	// Wait briefly for workflow to initialize
+	time.Sleep(time.Millisecond * 500)
+
+	var loadingPolicy *iwfidl.PersistenceLoadingPolicy
+	if useLocking {
+		// Use exclusive locking for synchronous updates
+		loadingPolicy = &iwfidl.PersistenceLoadingPolicy{
+			PersistenceLoadingType: iwfidl.PARTIAL_WITH_EXCLUSIVE_LOCK.Ptr(),
+			PartialLoadingKeys: []string{
+				rpcStorage.SmallDataKey,
+				rpcStorage.LargeDataKey,
+			},
+			LockingKeys: []string{
+				rpcStorage.SmallDataKey,
+				rpcStorage.LargeDataKey,
+			},
+		}
+	} else {
+		// Use non-locking for regular RPC
+		loadingPolicy = &iwfidl.PersistenceLoadingPolicy{
+			PersistenceLoadingType: iwfidl.ALL_WITHOUT_LOCKING.Ptr(),
+		}
+	}
+
+	// Test 1: Make RPC call to test external storage loading functionality
+	// The workflow waits for this RPC call to send an internal signal to close it
+	reqRpc := apiClient.DefaultApi.ApiV1WorkflowRpcPost(context.Background())
+	rpcResp, httpResp, err := reqRpc.WorkflowRpcRequest(iwfidl.WorkflowRpcRequest{
+		WorkflowId:                  wfId,
+		RpcName:                     rpcStorage.UpdateDataAttributesRPC,
+		Input:                       &rpcStorage.TestInput,
+		DataAttributesLoadingPolicy: loadingPolicy,
+		TimeoutSeconds:              iwfidl.PtrInt32(10),
+	}).Execute()
+
+	// The RPC should succeed since the workflow waits for it
+	failTestAtHttpError(err, httpResp, t)
+	assertions.Equal(&iwfidl.WorkflowRpcResponse{
+		Output: &rpcStorage.TestOutput,
+	}, rpcResp)
+	t.Logf("✅ RPC call succeeded and sent signal to close workflow")
+
+	// Give a moment for the worker handler to be called and store test data
+	time.Sleep(time.Millisecond * 100)
+
+	// Test 2: Verify the handler received correct data during RPC (this is the key test!)
+	_, testData := wfHandler.GetTestResult()
+
+	// Verify the RPC handler received the loaded data (not references)
+	rpcInputData, exists := testData[rpcStorage.UpdateDataAttributesRPC+"-received-data"]
+	assertions.True(exists, "RPC should have received data attributes")
+
+	receivedDataAttrs, ok := rpcInputData.([]iwfidl.KeyValue)
+	assertions.True(ok, "Received data should be KeyValue array")
+
+	// The handler should receive actual data content (loaded from external storage if needed)
+	receivedDataMap := make(map[string]string)
+	for _, attr := range receivedDataAttrs {
+		if attr.Value != nil && attr.Value.Data != nil {
+			receivedDataMap[*attr.Key] = *attr.Value.Data
+		}
+	}
+
+	// Both small and large data should be available as actual content to the handler
+	// This verifies that loadDataObjectsFromExternalStorage works correctly in RPC calls
+	if initialData, exists := receivedDataMap[rpcStorage.SmallDataKey]; exists {
+		assertions.Equal(*rpcStorage.InitialSmallData.Data, initialData, "Handler should receive initial small data content")
+	}
+	if initialLargeData, exists := receivedDataMap[rpcStorage.LargeDataKey]; exists {
+		assertions.Equal(*rpcStorage.InitialLargeData.Data, initialLargeData, "Handler should receive initial large data content (loaded from S3)")
+	}
+
+	t.Logf("✅ External storage functionality verified: RPC handler received actual data content, proving that large data was correctly loaded from external storage")
+
+	// The workflow should complete after the RPC sent the internal signal
+	// Wait for workflow to complete
+	reqWait := apiClient.DefaultApi.ApiV1WorkflowGetWithWaitPost(context.Background())
+	respWait, httpResp, err := reqWait.WorkflowGetRequest(iwfidl.WorkflowGetRequest{
+		WorkflowId:      wfId,
+		WaitTimeSeconds: iwfidl.PtrInt32(5),
+	}).Execute()
+	failTestAtHttpError(err, httpResp, t)
+
+	// Verify workflow completed successfully
+	assertions.Equal(iwfidl.COMPLETED, respWait.WorkflowStatus, "Workflow should complete successfully")
+}