chore: Introduce WithContext() for engine invocations (#590)

Author: PeterSchafer

pkg/workflow/engine_test.go

@@ -1,6 +1,7 @@
 package workflow
 
 import (
+	"context"
 	"fmt"
 	"net/url"
 	"testing"
@@ -524,3 +525,62 @@ func Test_EngineInvocationConcurrent(t *testing.T) {
 		}
 	}
 }
+
+func Test_EngineImpl_InvokeWithContext_CustomContext(t *testing.T) {
+	config := configuration.NewInMemory()
+	engine := NewWorkFlowEngine(config)
+
+	wfId := NewWorkflowIdentifier("ctxtest")
+	flagset := pflag.NewFlagSet("ctx", pflag.ContinueOnError)
+
+	var receivedCtx context.Context
+	_, err := engine.Register(wfId, ConfigurationOptionsFromFlagset(flagset), func(invocation InvocationContext, input []Data) ([]Data, error) {
+		receivedCtx = invocation.Context()
+		return nil, nil
+	})
+	assert.NoError(t, err)
+
+	err = engine.Init()
+	assert.NoError(t, err)
+
+	// Create a context with a specific value and deadline to verify it's passed through
+	type ctxKey string
+	testKey := ctxKey("test-key")
+	testValue := "test-value"
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	ctx = context.WithValue(ctx, testKey, testValue)
+
+	_, err = engine.Invoke(wfId, WithContext(ctx))
+	assert.NoError(t, err)
+	assert.NotNil(t, receivedCtx)
+	assert.Equal(t, testValue, receivedCtx.Value(testKey))
+
+	// Verify deadline is propagated
+	deadline, hasDeadline := receivedCtx.Deadline()
+	assert.True(t, hasDeadline, "context should have a deadline")
+	assert.False(t, deadline.IsZero(), "deadline should not be zero")
+}
+
+func Test_EngineImpl_InvokeWithContext_DefaultContext(t *testing.T) {
+	config := configuration.NewInMemory()
+	engine := NewWorkFlowEngine(config)
+
+	wfId := NewWorkflowIdentifier("ctxdefault")
+	flagset := pflag.NewFlagSet("cd", pflag.ContinueOnError)
+
+	var receivedCtx context.Context
+	_, err := engine.Register(wfId, ConfigurationOptionsFromFlagset(flagset), func(invocation InvocationContext, input []Data) ([]Data, error) {
+		receivedCtx = invocation.Context()
+		return nil, nil
+	})
+	assert.NoError(t, err)
+
+	err = engine.Init()
+	assert.NoError(t, err)
+
+	// Invoke without WithContext - should get a non-nil default context
+	_, err = engine.Invoke(wfId)
+	assert.NoError(t, err)
+	assert.NotNil(t, receivedCtx)
+}

pkg/workflow/engineimpl.go

@@ -1,6 +1,7 @@
 package workflow
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"net/url"
@@ -37,9 +38,10 @@ type EngineImpl struct {
 var _ Engine = (*EngineImpl)(nil)
 
 type engineRuntimeConfig struct {
-	config configuration.Configuration
-	input  []Data
-	ic     analytics.InstrumentationCollector
+	config  configuration.Configuration
+	input   []Data
+	ic      analytics.InstrumentationCollector
+	ctxFunc func() context.Context
 }
 
 type EngineInvokeOption func(*engineRuntimeConfig)
@@ -62,6 +64,14 @@ func WithInstrumentationCollector(ic analytics.InstrumentationCollector) EngineI
 	}
 }
 
+func WithContext(ctx context.Context) EngineInvokeOption {
+	return func(e *engineRuntimeConfig) {
+		if ctx != nil {
+			e.ctxFunc = func() context.Context { return ctx }
+		}
+	}
+}
+
 func (e *EngineImpl) GetLogger() *zerolog.Logger {
 	return e.logger
 }
@@ -321,11 +331,11 @@ func (e *EngineImpl) Invoke(
 			e.mu.Unlock()
 
 			// create a context object for the invocation
-			context := NewInvocationContext(id, options.config, localEngine, localNetworkAccess, localLogger, localAnalytics, localUi)
+			invocationCtx := newInvocationContext(options.ctxFunc, id, options.config, localEngine, localNetworkAccess, localLogger, localAnalytics, localUi)
 
 			// invoke workflow through its callback
 			localLogger.Printf("Workflow Start")
-			output, err = callback(context, options.input)
+			output, err = callback(invocationCtx, options.input)
 			localLogger.Printf("Workflow End")
 		}
 	} else {

pkg/workflow/invocationcontextimpl.go

@@ -5,6 +5,7 @@ import (
 	"log"
 
 	"github.com/rs/zerolog"
+
 	"github.com/snyk/go-application-framework/pkg/analytics"
 	"github.com/snyk/go-application-framework/pkg/configuration"
 	"github.com/snyk/go-application-framework/pkg/networking"
@@ -13,6 +14,7 @@ import (
 	"github.com/snyk/go-application-framework/pkg/utils"
 )
 
+// Deprecated: NewInvocationContext creates a new invocation context.
 func NewInvocationContext(
 	id Identifier,
 	config configuration.Configuration,
@@ -22,7 +24,24 @@ func NewInvocationContext(
 	analyticsImpl analytics.Analytics,
 	ui ui.UserInterface,
 ) InvocationContext {
+	return newInvocationContext(nil, id, config, engine, network, logger, analyticsImpl, ui)
+}
+
+func newInvocationContext(
+	ctxFunc func() context.Context,
+	id Identifier,
+	config configuration.Configuration,
+	engine Engine,
+	network networking.NetworkAccess,
+	logger zerolog.Logger,
+	analyticsImpl analytics.Analytics,
+	ui ui.UserInterface,
+) InvocationContext {
+	if ctxFunc == nil {
+		ctxFunc = context.Background
+	}
 	return &invocationContextImpl{
+		ctxFunc:        ctxFunc,
 		WorkflowID:     id,
 		Configuration:  config,
 		WorkflowEngine: engine,
@@ -35,6 +54,7 @@ func NewInvocationContext(
 
 // invocationContextImpl is the default implementation of the InvocationContext interface.
 type invocationContextImpl struct {
+	ctxFunc        func() context.Context
 	WorkflowID     Identifier
 	WorkflowEngine Engine
 	Configuration  configuration.Configuration
@@ -47,10 +67,8 @@ type invocationContextImpl struct {
 var _ InvocationContext = (*invocationContextImpl)(nil)
 
 // Context returns the context of the workflow that is being invoked.
-func (*invocationContextImpl) Context() context.Context {
-	// TODO: This is using context.Background() as a placeholder. Ideally this returns
-	// the context representing the lifecycle of the workflow that is being invoked.
-	return context.Background()
+func (ici *invocationContextImpl) Context() context.Context {
+	return ici.ctxFunc()
 }
 
 // GetWorkflowIdentifier returns the identifier of the workflow that is being invoked.