diff --git a/pkg/capabilities/capabilities.go b/pkg/capabilities/capabilities.go index ce5074678..7bc081d43 100644 --- a/pkg/capabilities/capabilities.go +++ b/pkg/capabilities/capabilities.go @@ -117,10 +117,11 @@ type RequestMetadata struct { } func (m *RequestMetadata) ContextWithCRE(ctx context.Context) context.Context { - return contexts.WithCRE(ctx, contexts.CRE{ - Owner: m.WorkflowOwner, - Workflow: m.WorkflowID, - }) + val := contexts.CREValue(ctx) + // preserve org, if set + val.Owner = m.WorkflowOwner + val.Workflow = m.WorkflowID + return contexts.WithCRE(ctx, val) } type RegistrationMetadata struct { @@ -131,10 +132,11 @@ type RegistrationMetadata struct { } func (m *RegistrationMetadata) ContextWithCRE(ctx context.Context) context.Context { - return contexts.WithCRE(ctx, contexts.CRE{ - Owner: m.WorkflowOwner, - Workflow: m.WorkflowID, - }) + val := contexts.CREValue(ctx) + // preserve org, if set + val.Owner = m.WorkflowOwner + val.Workflow = m.WorkflowID + return contexts.WithCRE(ctx, val) } // CapabilityRequest is a struct for the Execute request of a capability. diff --git a/pkg/capabilities/capabilities_test.go b/pkg/capabilities/capabilities_test.go index 5060b0f17..06e3bac43 100644 --- a/pkg/capabilities/capabilities_test.go +++ b/pkg/capabilities/capabilities_test.go @@ -11,6 +11,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/smartcontractkit/chainlink-common/pkg/contexts" "github.com/smartcontractkit/chainlink-protos/cre/go/values" ) @@ -331,3 +332,31 @@ func TestChainSelectorLabel(t *testing.T) { } func ptr[T any](v T) *T { return &v } + +func TestRequestMetadata_ContextWithCRE(t *testing.T) { + ctx := t.Context() + require.Equal(t, "", contexts.CREValue(ctx).Org) + + // set it + ctx = contexts.WithCRE(ctx, contexts.CRE{Org: "org-id"}) + require.Equal(t, "org-id", contexts.CREValue(ctx).Org) + + // preserve it + md := RequestMetadata{WorkflowOwner: "owner-id", WorkflowID: "workflow-id"} + ctx = md.ContextWithCRE(ctx) + require.Equal(t, "org-id", contexts.CREValue(ctx).Org) +} + +func TestRegistrationMetadata_ContextWithCRE(t *testing.T) { + ctx := t.Context() + require.Equal(t, "", contexts.CREValue(ctx).Org) + + // set it + ctx = contexts.WithCRE(ctx, contexts.CRE{Org: "org-id"}) + require.Equal(t, "org-id", contexts.CREValue(ctx).Org) + + // preserve it + md := RegistrationMetadata{WorkflowOwner: "owner-id", WorkflowID: "workflow-id"} + ctx = md.ContextWithCRE(ctx) + require.Equal(t, "org-id", contexts.CREValue(ctx).Org) +}