github.com/devseccon/trivy@v0.47.1-0.20231123133102-bd902a0bd996/pkg/fanal/handler/handler_test.go (about)

     1  package handler_test
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"github.com/stretchr/testify/assert"
     8  	"github.com/stretchr/testify/require"
     9  
    10  	"github.com/devseccon/trivy/pkg/fanal/analyzer"
    11  	"github.com/devseccon/trivy/pkg/fanal/artifact"
    12  	"github.com/devseccon/trivy/pkg/fanal/handler"
    13  	"github.com/devseccon/trivy/pkg/fanal/types"
    14  )
    15  
    16  type fakeHook struct{}
    17  
    18  func (h fakeHook) Handle(ctx context.Context, result *analyzer.AnalysisResult, info *types.BlobInfo) error {
    19  	info.DiffID = "fake"
    20  	return nil
    21  }
    22  
    23  func (h fakeHook) Priority() int {
    24  	return 1
    25  }
    26  
    27  func (h fakeHook) Version() int { return 1 }
    28  
    29  func (h fakeHook) Type() types.HandlerType { return "fake" }
    30  
    31  func fakehookInit(_ artifact.Option) (handler.PostHandler, error) {
    32  	return fakeHook{}, nil
    33  }
    34  
    35  func TestManager_Versions(t *testing.T) {
    36  	tests := []struct {
    37  		name    string
    38  		disable []types.HandlerType
    39  		want    map[string]int
    40  	}{
    41  		{
    42  			name: "happy path",
    43  			want: map[string]int{
    44  				"fake": 1,
    45  			},
    46  		},
    47  		{
    48  			name:    "disable hooks",
    49  			disable: []types.HandlerType{"fake"},
    50  			want:    map[string]int{},
    51  		},
    52  	}
    53  
    54  	for _, tt := range tests {
    55  		t.Run(tt.name, func(t *testing.T) {
    56  			handler.RegisterPostHandlerInit("fake", fakehookInit)
    57  			defer handler.DeregisterPostHandler("fake")
    58  			m, err := handler.NewManager(artifact.Option{
    59  				DisabledHandlers: tt.disable,
    60  			})
    61  			require.NoError(t, err)
    62  			got := m.Versions()
    63  			assert.Equal(t, tt.want, got)
    64  		})
    65  	}
    66  }
    67  
    68  func TestManager_CallHooks(t *testing.T) {
    69  	tests := []struct {
    70  		name    string
    71  		disable []types.HandlerType
    72  		want    types.BlobInfo
    73  	}{
    74  		{
    75  			name: "happy path",
    76  			want: types.BlobInfo{
    77  				Digest: "digest",
    78  				DiffID: "fake",
    79  			},
    80  		},
    81  		{
    82  			name:    "disable hooks",
    83  			disable: []types.HandlerType{"fake"},
    84  			want: types.BlobInfo{
    85  				Digest: "digest",
    86  			},
    87  		},
    88  	}
    89  
    90  	for _, tt := range tests {
    91  		t.Run(tt.name, func(t *testing.T) {
    92  			handler.RegisterPostHandlerInit("fake", fakehookInit)
    93  			defer handler.DeregisterPostHandler("fake")
    94  			blob := types.BlobInfo{
    95  				Digest: "digest",
    96  			}
    97  			m, err := handler.NewManager(artifact.Option{
    98  				DisabledHandlers: tt.disable,
    99  			})
   100  			require.NoError(t, err)
   101  
   102  			err = m.PostHandle(context.TODO(), nil, &blob)
   103  			require.NoError(t, err)
   104  			assert.Equal(t, tt.want, blob)
   105  		})
   106  	}
   107  }