fix(trivy): consistent coverage for reset method + longer wait time between retries (#1272)

Signed-off-by: Ana-Roberta Lisca <ana.kagome@yahoo.com>

pkg/extensions/extension_search.go

@@ -71,12 +71,18 @@ func EnableSearchExtension(config *config.Config, storeController storage.StoreC
 }
 
 func downloadTrivyDB(interval time.Duration, sch *scheduler.Scheduler, cveInfo CveInfo, log log.Logger) {
-	generator := &trivyTaskGenerator{interval, cveInfo, log, pending, 0, time.Now(), &sync.Mutex{}}
+	generator := NewTrivyTaskGenerator(interval, cveInfo, log)
 
 	sch.SubmitGenerator(generator, interval, scheduler.HighPriority)
 }
 
-type trivyTaskGenerator struct {
+func NewTrivyTaskGenerator(interval time.Duration, cveInfo CveInfo, log log.Logger) *TrivyTaskGenerator {
+	generator := &TrivyTaskGenerator{interval, cveInfo, log, pending, 0, time.Now(), &sync.Mutex{}}
+
+	return generator
+}
+
+type TrivyTaskGenerator struct {
 	interval     time.Duration
 	cveInfo      CveInfo
 	log          log.Logger
@@ -86,7 +92,7 @@ type trivyTaskGenerator struct {
 	lock         *sync.Mutex
 }
 
-func (gen *trivyTaskGenerator) GenerateTask() (scheduler.Task, error) {
+func (gen *TrivyTaskGenerator) GenerateTask() (scheduler.Task, error) {
 	var newTask scheduler.Task
 
 	gen.lock.Lock()
@@ -100,7 +106,7 @@ func (gen *trivyTaskGenerator) GenerateTask() (scheduler.Task, error) {
 	return newTask, nil
 }
 
-func (gen *trivyTaskGenerator) IsDone() bool {
+func (gen *TrivyTaskGenerator) IsDone() bool {
 	gen.lock.Lock()
 	status := gen.status
 	gen.lock.Unlock()
@@ -108,7 +114,7 @@ func (gen *trivyTaskGenerator) IsDone() bool {
 	return status == done
 }
 
-func (gen *trivyTaskGenerator) Reset() {
+func (gen *TrivyTaskGenerator) Reset() {
 	gen.lock.Lock()
 	gen.status = pending
 	gen.waitTime = 0
@@ -118,12 +124,12 @@ func (gen *trivyTaskGenerator) Reset() {
 type trivyTask struct {
 	interval  time.Duration
 	cveInfo   cveinfo.CveInfo
-	generator *trivyTaskGenerator
+	generator *TrivyTaskGenerator
 	log       log.Logger
 }
 
 func newTrivyTask(interval time.Duration, cveInfo cveinfo.CveInfo,
-	generator *trivyTaskGenerator, log log.Logger,
+	generator *TrivyTaskGenerator, log log.Logger,
 ) *trivyTask {
 	return &trivyTask{interval, cveInfo, generator, log}
 }
@@ -135,7 +141,12 @@ func (trivyT *trivyTask) DoWork() error {
 	if err != nil {
 		trivyT.generator.lock.Lock()
 		trivyT.generator.status = pending
-		trivyT.generator.waitTime += time.Second
+
+		if trivyT.generator.waitTime == 0 {
+			trivyT.generator.waitTime = time.Second
+		}
+
+		trivyT.generator.waitTime *= 2
 		trivyT.generator.lastTaskTime = time.Now()
 		trivyT.generator.lock.Unlock()
 

pkg/extensions/extension_search_test.go

@@ -0,0 +1,65 @@
+//go:build search
+// +build search
+
+package extensions_test
+
+import (
+	"context"
+	"os"
+	"testing"
+	"time"
+
+	ispec "github.com/opencontainers/image-spec/specs-go/v1"
+	. "github.com/smartystreets/goconvey/convey"
+
+	. "zotregistry.io/zot/pkg/extensions"
+	cveinfo "zotregistry.io/zot/pkg/extensions/search/cve"
+	"zotregistry.io/zot/pkg/log"
+	"zotregistry.io/zot/pkg/meta/repodb"
+	"zotregistry.io/zot/pkg/scheduler"
+	"zotregistry.io/zot/pkg/storage"
+	. "zotregistry.io/zot/pkg/test"
+	"zotregistry.io/zot/pkg/test/mocks"
+)
+
+func TestTrivyDBGenerator(t *testing.T) {
+	Convey("Test trivy task scheduler reset", t, func() {
+		logFile, err := os.CreateTemp(t.TempDir(), "zot-log*.txt")
+		logPath := logFile.Name()
+		So(err, ShouldBeNil)
+
+		defer os.Remove(logFile.Name()) // clean up
+
+		logger := log.NewLogger("debug", logFile.Name())
+		sch := scheduler.NewScheduler(logger)
+
+		repoDB := &mocks.RepoDBMock{
+			GetRepoMetaFn: func(repo string) (repodb.RepoMetadata, error) {
+				return repodb.RepoMetadata{
+					Tags: map[string]repodb.Descriptor{
+						"tag": {MediaType: ispec.MediaTypeImageIndex},
+					},
+				}, nil
+			},
+		}
+		storeController := storage.StoreController{
+			DefaultStore: mocks.MockedImageStore{},
+		}
+
+		cveInfo := cveinfo.NewCVEInfo(storeController, repoDB, "", logger)
+		generator := NewTrivyTaskGenerator(time.Minute, cveInfo, logger)
+
+		sch.SubmitGenerator(generator, 12000*time.Millisecond, scheduler.HighPriority)
+
+		ctx, cancel := context.WithCancel(context.Background())
+		sch.RunScheduler(ctx)
+
+		defer cancel()
+
+		// Wait for trivy db to download
+		found, err := ReadLogFileAndCountStringOccurence(logPath,
+			"DB update completed, next update scheduled", 90*time.Second, 2)
+		So(err, ShouldBeNil)
+		So(found, ShouldBeTrue)
+	})
+}

pkg/test/common.go

@@ -1049,6 +1049,29 @@ func ReadLogFileAndSearchString(logPath string, stringToMatch string, timeout ti
 	}
 }
 
+func ReadLogFileAndCountStringOccurence(logPath string, stringToMatch string,
+	timeout time.Duration, count int,
+) (bool, error) {
+	ctx, cancelFunc := context.WithTimeout(context.Background(), timeout)
+	defer cancelFunc()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return false, nil
+		default:
+			content, err := os.ReadFile(logPath)
+			if err != nil {
+				return false, err
+			}
+
+			if strings.Count(string(content), stringToMatch) >= count {
+				return true, nil
+			}
+		}
+	}
+}
+
 func CopyFile(sourceFilePath, destFilePath string) error {
 	destFile, err := os.Create(destFilePath)
 	if err != nil {

pkg/test/common_test.go

@@ -725,6 +725,39 @@ func TestReadLogFileAndSearchString(t *testing.T) {
 	})
 }
 
+func TestReadLogFileAndCountStringOccurence(t *testing.T) {
+	logFile, err := os.CreateTemp(t.TempDir(), "zot-log*.txt")
+	if err != nil {
+		panic(err)
+	}
+
+	_, err = logFile.Write([]byte("line1\n line2\n line3 line1 line2\n line1"))
+	if err != nil {
+		panic(err)
+	}
+
+	logPath := logFile.Name()
+	defer os.Remove(logPath)
+
+	Convey("Invalid path", t, func() {
+		_, err = test.ReadLogFileAndCountStringOccurence("invalidPath",
+			"DB update completed, next update scheduled", 90*time.Second, 1)
+		So(err, ShouldNotBeNil)
+	})
+
+	Convey("Time too short", t, func() {
+		ok, err := test.ReadLogFileAndCountStringOccurence(logPath, "invalid string", time.Microsecond, 1)
+		So(err, ShouldBeNil)
+		So(ok, ShouldBeFalse)
+	})
+
+	Convey("Count occurrence working", t, func() {
+		ok, err := test.ReadLogFileAndCountStringOccurence(logPath, "line1", 90*time.Second, 3)
+		So(err, ShouldBeNil)
+		So(ok, ShouldBeTrue)
+	})
+}
+
 func TestInjectUploadImageWithBasicAuth(t *testing.T) {
 	Convey("Inject failures for unreachable lines", t, func() {
 		port := test.GetFreePort()