feat: context propagation on GetWorkflowSpec

smartcontractkit · Oct 25, 2024 · 2d1c349 · 2d1c349
1 parent a3f7620
commit 2d1c349
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 4 deletions.
diff --git a/core/services/job/wasm_file_spec_factory.go b/core/services/job/wasm_file_spec_factory.go
@@ -21,7 +21,7 @@ import (
 
 type WasmFileSpecFactory struct{}
 
-func (w WasmFileSpecFactory) Spec(_ context.Context, workflow, configLocation string) (sdk.WorkflowSpec, []byte, string, error) {
+func (w WasmFileSpecFactory) Spec(ctx context.Context, workflow, configLocation string) (sdk.WorkflowSpec, []byte, string, error) {
 	config, err := os.ReadFile(configLocation)
 	if err != nil {
 		return sdk.WorkflowSpec{}, nil, "", err
@@ -33,7 +33,7 @@ func (w WasmFileSpecFactory) Spec(_ context.Context, workflow, configLocation st
 	}
 
 	moduleConfig := &host.ModuleConfig{Logger: logger.NullLogger}
-	spec, err := host.GetWorkflowSpec(moduleConfig, compressedBinary, config)
+	spec, err := host.GetWorkflowSpec(ctx, moduleConfig, compressedBinary, config)
 	if err != nil {
 		return sdk.WorkflowSpec{}, nil, "", err
 	} else if spec == nil {

diff --git a/core/services/job/wasm_file_spec_factory_test.go b/core/services/job/wasm_file_spec_factory_test.go
@@ -37,11 +37,12 @@ func TestWasmFileSpecFactory(t *testing.T) {
 	require.NoError(t, bwr.Close())
 
 	t.Run("Raw binary", func(t *testing.T) {
+		ctx := testutils.Context(t)
 		factory := job.WasmFileSpecFactory{}
 		actual, rawSpec, actualSha, err2 := factory.Spec(testutils.Context(t), binaryLocation, configLocation)
 		require.NoError(t, err2)
 
-		expected, err2 := host.GetWorkflowSpec(&host.ModuleConfig{Logger: logger.NullLogger, IsUncompressed: true}, rawBinary, config)
+		expected, err2 := host.GetWorkflowSpec(ctx, &host.ModuleConfig{Logger: logger.NullLogger, IsUncompressed: true}, rawBinary, config)
 		require.NoError(t, err2)
 
 		expectedSha := sha256.New()
@@ -55,6 +56,7 @@ func TestWasmFileSpecFactory(t *testing.T) {
 	})
 
 	t.Run("Compressed binary", func(t *testing.T) {
+		ctx := testutils.Context(t)
 		brLoc := strings.Replace(binaryLocation, ".wasm", ".br", 1)
 		compressedBytes := b.Bytes()
 		require.NoError(t, os.WriteFile(brLoc, compressedBytes, 0600))
@@ -63,7 +65,7 @@ func TestWasmFileSpecFactory(t *testing.T) {
 		actual, rawSpec, actualSha, err2 := factory.Spec(testutils.Context(t), brLoc, configLocation)
 		require.NoError(t, err2)
 
-		expected, err2 := host.GetWorkflowSpec(&host.ModuleConfig{Logger: logger.NullLogger, IsUncompressed: true}, rawBinary, config)
+		expected, err2 := host.GetWorkflowSpec(ctx, &host.ModuleConfig{Logger: logger.NullLogger, IsUncompressed: true}, rawBinary, config)
 		require.NoError(t, err2)
 
 		expectedSha := sha256.New()

diff --git a/core/services/workflows/engine_test.go b/core/services/workflows/engine_test.go
@@ -1449,6 +1449,7 @@ func TestEngine_WithCustomComputeStep(t *testing.T) {
 	binaryB := wasmtest.CreateTestBinary(cmd, binary, true, t)
 
 	spec, err := host.GetWorkflowSpec(
+		ctx,
 		&host.ModuleConfig{Logger: log},
 		binaryB,
 		nil, // config
@@ -1511,6 +1512,7 @@ func TestEngine_CustomComputePropagatesBreaks(t *testing.T) {
 	binaryB := wasmtest.CreateTestBinary(cmd, binary, true, t)
 
 	spec, err := host.GetWorkflowSpec(
+		ctx,
 		&host.ModuleConfig{Logger: log},
 		binaryB,
 		nil, // config