github.com/devseccon/trivy@v0.47.1-0.20231123133102-bd902a0bd996/pkg/scanner/post/post_scan_test.go (about)

     1  package post_test
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"testing"
     7  
     8  	"github.com/stretchr/testify/assert"
     9  	"github.com/stretchr/testify/require"
    10  
    11  	dbTypes "github.com/aquasecurity/trivy-db/pkg/types"
    12  	"github.com/devseccon/trivy/pkg/scanner/post"
    13  	"github.com/devseccon/trivy/pkg/types"
    14  )
    15  
    16  type testPostScanner struct{}
    17  
    18  func (testPostScanner) Name() string {
    19  	return "test"
    20  }
    21  
    22  func (testPostScanner) Version() int {
    23  	return 1
    24  }
    25  
    26  func (testPostScanner) PostScan(ctx context.Context, results types.Results) (types.Results, error) {
    27  	for i, r := range results {
    28  		if r.Target == "bad" {
    29  			return nil, errors.New("bad")
    30  		}
    31  		for j := range r.Vulnerabilities {
    32  			results[i].Vulnerabilities[j].Severity = "LOW"
    33  		}
    34  	}
    35  	return results, nil
    36  }
    37  
    38  func TestScan(t *testing.T) {
    39  	tests := []struct {
    40  		name    string
    41  		results types.Results
    42  		want    types.Results
    43  		wantErr bool
    44  	}{
    45  		{
    46  			name: "happy path",
    47  			results: types.Results{
    48  				{
    49  					Target: "test",
    50  					Vulnerabilities: []types.DetectedVulnerability{
    51  						{
    52  							VulnerabilityID:  "CVE-2022-0001",
    53  							PkgName:          "musl",
    54  							InstalledVersion: "1.2.3",
    55  							FixedVersion:     "1.2.4",
    56  							Vulnerability: dbTypes.Vulnerability{
    57  								Severity: "CRITICAL",
    58  							},
    59  						},
    60  					},
    61  				},
    62  			},
    63  			want: types.Results{
    64  				{
    65  					Target: "test",
    66  					Vulnerabilities: []types.DetectedVulnerability{
    67  						{
    68  							VulnerabilityID:  "CVE-2022-0001",
    69  							PkgName:          "musl",
    70  							InstalledVersion: "1.2.3",
    71  							FixedVersion:     "1.2.4",
    72  							Vulnerability: dbTypes.Vulnerability{
    73  								Severity: "LOW",
    74  							},
    75  						},
    76  					},
    77  				},
    78  			},
    79  		},
    80  		{
    81  			name: "sad path",
    82  			results: types.Results{
    83  				{
    84  					Target: "bad",
    85  				},
    86  			},
    87  			wantErr: true,
    88  		},
    89  	}
    90  	for _, tt := range tests {
    91  		t.Run(tt.name, func(t *testing.T) {
    92  			s := testPostScanner{}
    93  			post.RegisterPostScanner(s)
    94  			defer func() {
    95  				post.DeregisterPostScanner(s.Name())
    96  			}()
    97  
    98  			results, err := post.Scan(context.Background(), tt.results)
    99  			require.Equal(t, err != nil, tt.wantErr)
   100  			assert.Equal(t, results, tt.want)
   101  		})
   102  	}
   103  }