Skip to content

Commit

Permalink
Updating component runners to leverage WorkflowID.
Browse files Browse the repository at this point in the history
  • Loading branch information
@andream16
andream16 committed Oct 30, 2024
1 parent f33586e commit 9e6e9db
Show file tree
Hide file tree
Showing 9 changed files with 138 additions and 97 deletions.
9 changes: 5 additions & 4 deletions sdk/component/enricher.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,9 @@ func RunEnricher(ctx context.Context, enricher Enricher, opts ...RunnerOption) e
ctx,
func(ctx context.Context, cfg *RunnerConfig) error {
var (
logger = LoggerFromContext(ctx).With(logKeyComponentType, "enricher")
store = cfg.storerConfig.store
workflowID = cfg.WorkflowID
logger = LoggerFromContext(ctx).With(logKeyComponentType, "enricher")
store = cfg.storerConfig.store
)

defer func() {
Expand All @@ -24,7 +25,7 @@ func RunEnricher(ctx context.Context, enricher Enricher, opts ...RunnerOption) e
logger.Debug("preparing to execute enricher component...")
logger.Debug("preparing to execute read step...")

findings, err := store.Read(ctx)
findings, err := store.Read(ctx, workflowID)
if err != nil {
logger.With(logKeyError, err.Error()).Error("reading step failed")
return fmt.Errorf("could not read: %w", err)
Expand All @@ -44,7 +45,7 @@ func RunEnricher(ctx context.Context, enricher Enricher, opts ...RunnerOption) e
logger.Debug("enricher step completed!")
logger.Debug("preparing to execute update step...")

if err := store.Update(ctx, enrichedFindings); err != nil {
if err := store.Update(ctx, workflowID, enrichedFindings); err != nil {
logger.With(logKeyError, err.Error()).Error("updating step failed")
return fmt.Errorf("could not update: %w", err)
}
Expand Down
65 changes: 38 additions & 27 deletions sdk/component/enricher_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (
"go.uber.org/mock/gomock"

"github.com/smithy-security/smithy/sdk/component/internal/mocks"
"github.com/smithy-security/smithy/sdk/component/internal/uuid"

"github.com/smithy-security/smithy/sdk/component"
ocsf "github.com/smithy-security/smithy/sdk/gen/com/github/ocsf/ocsf_schema/v1"
Expand All @@ -18,6 +19,7 @@ import (
func runEnricherHelper(
t *testing.T,
ctx context.Context,
workflowID uuid.UUID,
enricher component.Enricher,
store component.Storer,
) error {
Expand All @@ -28,13 +30,15 @@ func runEnricherHelper(
enricher,
component.RunnerWithLogger(component.NewNoopLogger()),
component.RunnerWithComponentName("sample-enricher"),
component.RunnerWithWorkflowID(workflowID),
component.RunnerWithStorer("local", store),
)
}

func TestRunEnricher(t *testing.T) {
var (
ctrl, ctx = gomock.WithContext(context.Background(), t)
workflowID = uuid.New()
mockCtx = gomock.AssignableToTypeOf(ctx)
mockStore = mocks.NewMockStorer(ctrl)
mockEnricher = mocks.NewMockEnricher(ctrl)
Expand All @@ -46,23 +50,23 @@ func TestRunEnricher(t *testing.T) {
gomock.InOrder(
mockStore.
EXPECT().
Read(mockCtx).
Read(mockCtx, workflowID).
Return(vulns, nil),
mockEnricher.
EXPECT().
Annotate(mockCtx, vulns).
Return(enrichedVulns, nil),
mockStore.
EXPECT().
Update(mockCtx, enrichedVulns).
Update(mockCtx, workflowID, enrichedVulns).
Return(nil),
mockStore.
EXPECT().
Close(mockCtx).
Return(nil),
)

require.NoError(t, runEnricherHelper(t, ctx, mockEnricher, mockStore))
require.NoError(t, runEnricherHelper(t, ctx, workflowID, mockEnricher, mockStore))
})

t.Run("it should return early when the context is cancelled", func(t *testing.T) {
Expand All @@ -71,29 +75,35 @@ func TestRunEnricher(t *testing.T) {
gomock.InOrder(
mockStore.
EXPECT().
Read(mockCtx).
Read(mockCtx, workflowID).
Return(vulns, nil),
mockEnricher.
EXPECT().
Annotate(mockCtx, vulns).
DoAndReturn(func(ctx context.Context, vulns []*ocsf.VulnerabilityFinding) ([]*ocsf.VulnerabilityFinding, error) {
cancel()
return enrichedVulns, nil
}),
DoAndReturn(
func(ctx context.Context, vulns []*ocsf.VulnerabilityFinding) ([]*ocsf.VulnerabilityFinding, error) {
cancel()
return enrichedVulns, nil
}),
mockStore.
EXPECT().
Update(mockCtx, enrichedVulns).
DoAndReturn(func(ctx context.Context, vulns []*ocsf.VulnerabilityFinding) error {
<-ctx.Done()
return nil
}),
Update(mockCtx, workflowID, enrichedVulns).
DoAndReturn(
func(
ctx context.Context,
workflowID uuid.UUID,
vulns []*ocsf.VulnerabilityFinding,
) error {
<-ctx.Done()
return nil
}),
mockStore.
EXPECT().
Close(mockCtx).
Return(nil),
)

require.NoError(t, runEnricherHelper(t, ctx, mockEnricher, mockStore))
require.NoError(t, runEnricherHelper(t, ctx, workflowID, mockEnricher, mockStore))
})

t.Run("it should return early when reading errors", func(t *testing.T) {
Expand All @@ -102,15 +112,15 @@ func TestRunEnricher(t *testing.T) {
gomock.InOrder(
mockStore.
EXPECT().
Read(mockCtx).
Read(mockCtx, workflowID).
Return(nil, errRead),
mockStore.
EXPECT().
Close(mockCtx).
Return(nil),
)

err := runEnricherHelper(t, ctx, mockEnricher, mockStore)
err := runEnricherHelper(t, ctx, workflowID, mockEnricher, mockStore)
require.Error(t, err)
assert.ErrorIs(t, err, errRead)
})
Expand All @@ -121,7 +131,7 @@ func TestRunEnricher(t *testing.T) {
gomock.InOrder(
mockStore.
EXPECT().
Read(mockCtx).
Read(mockCtx, workflowID).
Return(vulns, nil),
mockEnricher.
EXPECT().
Expand All @@ -133,7 +143,7 @@ func TestRunEnricher(t *testing.T) {
Return(nil),
)

err := runEnricherHelper(t, ctx, mockEnricher, mockStore)
err := runEnricherHelper(t, ctx, workflowID, mockEnricher, mockStore)
require.Error(t, err)
assert.ErrorIs(t, err, errAnnotation)
})
Expand All @@ -144,23 +154,23 @@ func TestRunEnricher(t *testing.T) {
gomock.InOrder(
mockStore.
EXPECT().
Read(mockCtx).
Read(mockCtx, workflowID).
Return(vulns, nil),
mockEnricher.
EXPECT().
Annotate(mockCtx, vulns).
Return(enrichedVulns, nil),
mockStore.
EXPECT().
Update(mockCtx, enrichedVulns).
Update(mockCtx, workflowID, enrichedVulns).
Return(errUpdate),
mockStore.
EXPECT().
Close(mockCtx).
Return(nil),
)

err := runEnricherHelper(t, ctx, mockEnricher, mockStore)
err := runEnricherHelper(t, ctx, workflowID, mockEnricher, mockStore)
require.Error(t, err)
assert.ErrorIs(t, err, errUpdate)
})
Expand All @@ -171,22 +181,23 @@ func TestRunEnricher(t *testing.T) {
gomock.InOrder(
mockStore.
EXPECT().
Read(mockCtx).
Read(mockCtx, workflowID).
Return(vulns, nil),
mockEnricher.
EXPECT().
Annotate(mockCtx, vulns).
DoAndReturn(func(ctx context.Context, vulns []*ocsf.VulnerabilityFinding) ([]*ocsf.VulnerabilityFinding, error) {
panic(errAnnotation)
return enrichedVulns, nil
}),
DoAndReturn(
func(ctx context.Context, vulns []*ocsf.VulnerabilityFinding) ([]*ocsf.VulnerabilityFinding, error) {
panic(errAnnotation)
return enrichedVulns, nil
}),
mockStore.
EXPECT().
Close(mockCtx).
Return(nil),
)

err := runEnricherHelper(t, ctx, mockEnricher, mockStore)
err := runEnricherHelper(t, ctx, workflowID, mockEnricher, mockStore)
require.Error(t, err)
assert.ErrorIs(t, err, errAnnotation)
})
Expand Down
9 changes: 5 additions & 4 deletions sdk/component/filter.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,9 @@ func RunFilter(ctx context.Context, filter Filter, opts ...RunnerOption) error {
ctx,
func(ctx context.Context, cfg *RunnerConfig) error {
var (
logger = LoggerFromContext(ctx).With(logKeyComponentType, "filter")
store = cfg.storerConfig.store
workflowID = cfg.WorkflowID
logger = LoggerFromContext(ctx).With(logKeyComponentType, "filter")
store = cfg.storerConfig.store
)

defer func() {
Expand All @@ -24,7 +25,7 @@ func RunFilter(ctx context.Context, filter Filter, opts ...RunnerOption) error {
logger.Debug("preparing to execute filter component...")
logger.Debug("preparing to execute read step...")

findings, err := store.Read(ctx)
findings, err := store.Read(ctx, workflowID)
if err != nil {
logger.With(logKeyError, err.Error()).Error("reading step failed")
return fmt.Errorf("could not read: %w", err)
Expand All @@ -48,7 +49,7 @@ func RunFilter(ctx context.Context, filter Filter, opts ...RunnerOption) error {
logger.Debug("filter step completed!")
logger.Debug("preparing to execute update step...")

if err := store.Update(ctx, filteredFindings); err != nil {
if err := store.Update(ctx, workflowID, filteredFindings); err != nil {
logger.With(logKeyError, err.Error()).Error("updating step failed")
return fmt.Errorf("could not update: %w", err)
}
Expand Down
Loading

0 comments on commit 9e6e9db

Please sign in to comment.