zotregistry.io/zot@v1.4.4-0.20231124084042-02a8ed785457/pkg/extensions/search/cve/update_test.go (about)

     1  //go:build search
     2  // +build search
     3  
     4  package cveinfo_test
     5  
     6  import (
     7  	"context"
     8  	"io"
     9  	"os"
    10  	"testing"
    11  	"time"
    12  
    13  	ispec "github.com/opencontainers/image-spec/specs-go/v1"
    14  	. "github.com/smartystreets/goconvey/convey"
    15  
    16  	"zotregistry.io/zot/pkg/api/config"
    17  	cveinfo "zotregistry.io/zot/pkg/extensions/search/cve"
    18  	"zotregistry.io/zot/pkg/log"
    19  	mTypes "zotregistry.io/zot/pkg/meta/types"
    20  	"zotregistry.io/zot/pkg/scheduler"
    21  	"zotregistry.io/zot/pkg/storage"
    22  	test "zotregistry.io/zot/pkg/test/common"
    23  	"zotregistry.io/zot/pkg/test/mocks"
    24  )
    25  
    26  func TestCVEDBGenerator(t *testing.T) {
    27  	Convey("Test CVE DB task scheduler reset", t, func() {
    28  		logFile, err := os.CreateTemp(t.TempDir(), "zot-log*.txt")
    29  		logPath := logFile.Name()
    30  		So(err, ShouldBeNil)
    31  
    32  		defer os.Remove(logFile.Name()) // clean up
    33  
    34  		logger := log.NewLogger("debug", logPath)
    35  		writers := io.MultiWriter(os.Stdout, logFile)
    36  		logger.Logger = logger.Output(writers)
    37  
    38  		cfg := config.New()
    39  		cfg.Scheduler = &config.SchedulerConfig{NumWorkers: 3}
    40  		sch := scheduler.NewScheduler(cfg, logger)
    41  
    42  		metaDB := &mocks.MetaDBMock{
    43  			GetRepoMetaFn: func(ctx context.Context, repo string) (mTypes.RepoMeta, error) {
    44  				return mTypes.RepoMeta{
    45  					Tags: map[string]mTypes.Descriptor{
    46  						"tag": {MediaType: ispec.MediaTypeImageIndex},
    47  					},
    48  				}, nil
    49  			},
    50  		}
    51  		storeController := storage.StoreController{
    52  			DefaultStore: mocks.MockedImageStore{
    53  				RootDirFn: func() string {
    54  					return t.TempDir()
    55  				},
    56  			},
    57  		}
    58  
    59  		cveScanner := cveinfo.NewScanner(storeController, metaDB, "ghcr.io/project-zot/trivy-db", "", logger)
    60  		generator := cveinfo.NewDBUpdateTaskGenerator(time.Minute, cveScanner, logger)
    61  
    62  		sch.SubmitGenerator(generator, 12000*time.Millisecond, scheduler.HighPriority)
    63  
    64  		ctx, cancel := context.WithCancel(context.Background())
    65  
    66  		sch.RunScheduler(ctx)
    67  
    68  		defer cancel()
    69  
    70  		// Wait for trivy db to download
    71  		found, err := test.ReadLogFileAndCountStringOccurence(logPath,
    72  			"DB update completed, next update scheduled", 140*time.Second, 2)
    73  		So(err, ShouldBeNil)
    74  		So(found, ShouldBeTrue)
    75  	})
    76  }