zotregistry.dev/zot@v1.4.4-0.20240314164342-eec277e14d20/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.dev/zot/pkg/api/config"
    17  	"zotregistry.dev/zot/pkg/extensions/monitoring"
    18  	cveinfo "zotregistry.dev/zot/pkg/extensions/search/cve"
    19  	"zotregistry.dev/zot/pkg/log"
    20  	mTypes "zotregistry.dev/zot/pkg/meta/types"
    21  	"zotregistry.dev/zot/pkg/scheduler"
    22  	"zotregistry.dev/zot/pkg/storage"
    23  	test "zotregistry.dev/zot/pkg/test/common"
    24  	"zotregistry.dev/zot/pkg/test/mocks"
    25  )
    26  
    27  func TestCVEDBGenerator(t *testing.T) {
    28  	Convey("Test CVE DB task scheduler reset", t, func() {
    29  		logFile, err := os.CreateTemp(t.TempDir(), "zot-log*.txt")
    30  		logPath := logFile.Name()
    31  		So(err, ShouldBeNil)
    32  
    33  		defer os.Remove(logFile.Name()) // clean up
    34  
    35  		logger := log.NewLogger("debug", logPath)
    36  		writers := io.MultiWriter(os.Stdout, logFile)
    37  		logger.Logger = logger.Output(writers)
    38  
    39  		cfg := config.New()
    40  		cfg.Scheduler = &config.SchedulerConfig{NumWorkers: 3}
    41  		metrics := monitoring.NewMetricsServer(true, logger)
    42  		sch := scheduler.NewScheduler(cfg, metrics, logger)
    43  
    44  		metaDB := &mocks.MetaDBMock{
    45  			GetRepoMetaFn: func(ctx context.Context, repo string) (mTypes.RepoMeta, error) {
    46  				return mTypes.RepoMeta{
    47  					Tags: map[mTypes.Tag]mTypes.Descriptor{
    48  						"tag": {MediaType: ispec.MediaTypeImageIndex},
    49  					},
    50  				}, nil
    51  			},
    52  		}
    53  		storeController := storage.StoreController{
    54  			DefaultStore: mocks.MockedImageStore{
    55  				RootDirFn: func() string {
    56  					return t.TempDir()
    57  				},
    58  			},
    59  		}
    60  
    61  		cveScanner := cveinfo.NewScanner(storeController, metaDB, "ghcr.io/project-zot/trivy-db", "", logger)
    62  		generator := cveinfo.NewDBUpdateTaskGenerator(time.Minute, cveScanner, logger)
    63  
    64  		sch.SubmitGenerator(generator, 12000*time.Millisecond, scheduler.HighPriority)
    65  
    66  		sch.RunScheduler()
    67  
    68  		defer sch.Shutdown()
    69  
    70  		// Wait for trivy db to download
    71  		found, err := test.ReadLogFileAndCountStringOccurence(logPath,
    72  			"cve-db update completed, next update scheduled after interval", 140*time.Second, 2)
    73  		So(err, ShouldBeNil)
    74  		So(found, ShouldBeTrue)
    75  	})
    76  }