diff --git a/dbos/dbos.go b/dbos/dbos.go index f9aa829..2c080a5 100644 --- a/dbos/dbos.go +++ b/dbos/dbos.go @@ -361,6 +361,10 @@ func NewDBOSContext(ctx context.Context, inputConfig Config) (DBOSContext, error safeGobRegister(ws, initExecutor.logger) var si []StepInfo safeGobRegister(si, initExecutor.logger) + var h workflowHandle[any] + safeGobRegister(h, initExecutor.logger) + var ph workflowPollingHandle[any] + safeGobRegister(ph, initExecutor.logger) // Initialize global variables from processed config (already handles env vars and defaults) initExecutor.applicationVersion = config.ApplicationVersion diff --git a/dbos/dbos_test.go b/dbos/dbos_test.go index 43dbcee..2840335 100644 --- a/dbos/dbos_test.go +++ b/dbos/dbos_test.go @@ -719,3 +719,54 @@ func TestCustomPool(t *testing.T) { assert.False(t, systemDB.(*sysDB).launched) }) } + +func TestWorkflowHandlerRegistration(t *testing.T) { + databaseURL := getDatabaseURL() + ctx, err := NewDBOSContext(context.Background(), Config{ + DatabaseURL: databaseURL, + AppName: "test-initialize", + }) + require.NoError(t, err) + defer func() { + if ctx != nil { + Shutdown(ctx, 1*time.Minute) + } + }() // Clean up executor + + require.NotNil(t, ctx) + + childWf := func(ctx DBOSContext, input string) (string, error) { + return "child-" + input, nil + } + + wf := func(ctx DBOSContext, input string) (WorkflowHandle[string], error) { + childHandle, err := RunWorkflow(ctx, childWf, input) + if err != nil { + var emptyHandle WorkflowHandle[string] + return emptyHandle, err + } + return childHandle, nil + } + + RegisterWorkflow(ctx, childWf) + RegisterWorkflow(ctx, wf) + + // Launch a workflow + err = Launch(ctx) + require.NoError(t, err) + defer Shutdown(ctx, 1*time.Minute) + + parentHandle, err := RunWorkflow(ctx, wf, "test-input") + require.NoError(t, err) + + // The result of the parent workflow should itself be a workflow handle + childHandle, err := parentHandle.GetResult() + require.NoError(t, err, "failed to get child handle from parent workflow") + + // Now we can use the child handle to get the final result + result, err := childHandle.GetResult() + require.NoError(t, err, "failed to get result from child workflow") + + // Assert final child workflow result + assert.Equal(t, "child-test-input", result) +}