chore: fix tests and add tests

Author: PeterSchafer

pkg/local_workflows/report_analytics_workflow_test.go

@@ -2,6 +2,7 @@ package localworkflows
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"io"
 	"net/http"
@@ -46,6 +47,7 @@ func Test_ReportAnalytics_ReportAnalyticsEntryPoint_shouldReportV2AnalyticsPaylo
 	invocationContextMock.EXPECT().GetEnhancedLogger().Return(&logger).AnyTimes()
 	invocationContextMock.EXPECT().GetEngine().Return(engineMock).AnyTimes()
 	invocationContextMock.EXPECT().GetNetworkAccess().Return(networkAccessMock).AnyTimes()
+	invocationContextMock.EXPECT().Context().Return(context.Background()).AnyTimes()
 	networkAccessMock.EXPECT().GetHttpClient().Return(mockClient).AnyTimes()
 
 	_, err := reportAnalyticsEntrypoint(invocationContextMock, []workflow.Data{testPayload(requestPayload)})
@@ -84,6 +86,7 @@ func Test_ReportAnalytics_ReportAnalyticsEntryPoint_reportsHttpStatusError(t *te
 	invocationContextMock.EXPECT().GetConfiguration().Return(config).AnyTimes()
 	invocationContextMock.EXPECT().GetEnhancedLogger().Return(&logger).AnyTimes()
 	invocationContextMock.EXPECT().GetNetworkAccess().Return(networkAccessMock).AnyTimes()
+	invocationContextMock.EXPECT().Context().Return(context.Background()).AnyTimes()
 	networkAccessMock.EXPECT().GetHttpClient().Return(mockClient).AnyTimes()
 
 	_, err := reportAnalyticsEntrypoint(invocationContextMock, []workflow.Data{testPayload(requestPayload)})
@@ -112,6 +115,7 @@ func Test_ReportAnalytics_ReportAnalyticsEntryPoint_reportsHttpError(t *testing.
 	invocationContextMock.EXPECT().GetConfiguration().Return(config).AnyTimes()
 	invocationContextMock.EXPECT().GetEnhancedLogger().Return(&logger).AnyTimes()
 	invocationContextMock.EXPECT().GetNetworkAccess().Return(networkAccessMock).AnyTimes()
+	invocationContextMock.EXPECT().Context().Return(context.Background()).AnyTimes()
 	networkAccessMock.EXPECT().GetHttpClient().Return(mockClient).AnyTimes()
 
 	_, err := reportAnalyticsEntrypoint(invocationContextMock, []workflow.Data{testPayload(requestPayload)})
@@ -172,6 +176,7 @@ func Test_ReportAnalytics_ReportAnalyticsEntryPoint_usesCLIInput(t *testing.T) {
 	invocationContextMock.EXPECT().GetNetworkAccess().Return(networkAccessMock).AnyTimes()
 	invocationContextMock.EXPECT().GetEngine().Return(engineMock).AnyTimes()
 	invocationContextMock.EXPECT().GetRuntimeInfo().Return(runtimeinfo.New(runtimeinfo.WithName("snyk-cli"), runtimeinfo.WithVersion("1.1233.0"))).AnyTimes()
+	invocationContextMock.EXPECT().Context().Return(context.Background()).AnyTimes()
 	engineMock.EXPECT().GetWorkflows().AnyTimes()
 	networkAccessMock.EXPECT().GetHttpClient().Return(mockClient).AnyTimes()
 

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)
+}