refactor: delete db mock

integration/integration_test.go

@@ -28,9 +28,10 @@ import (
 
 	"github.com/aquasecurity/trivy-db/pkg/db"
 	"github.com/aquasecurity/trivy-db/pkg/metadata"
+
+	"github.com/aquasecurity/trivy/internal/dbtest"
 	"github.com/aquasecurity/trivy/pkg/clock"
 	"github.com/aquasecurity/trivy/pkg/commands"
-	"github.com/aquasecurity/trivy/pkg/dbtest"
 	"github.com/aquasecurity/trivy/pkg/types"
 	"github.com/aquasecurity/trivy/pkg/uuid"
 

internal/dbtest/fake.go

@@ -0,0 +1,84 @@
+package dbtest
+
+import (
+	"archive/tar"
+	"os"
+	"path/filepath"
+	"testing"
+
+	v1 "github.com/google/go-containerregistry/pkg/v1"
+	fakei "github.com/google/go-containerregistry/pkg/v1/fake"
+	"github.com/google/go-containerregistry/pkg/v1/tarball"
+	"github.com/google/go-containerregistry/pkg/v1/types"
+	"github.com/samber/lo"
+	"github.com/stretchr/testify/require"
+
+	ftypes "github.com/aquasecurity/trivy/pkg/fanal/types"
+	"github.com/aquasecurity/trivy/pkg/oci"
+)
+
+const defaultMediaType = "application/vnd.aquasec.trivy.db.layer.v1.tar+gzip"
+
+type fakeLayer struct {
+	v1.Layer
+}
+
+func (f fakeLayer) MediaType() (types.MediaType, error) {
+	return f.Layer.MediaType()
+}
+
+func NewFakeLayer(t *testing.T, input string, mediaType types.MediaType) v1.Layer {
+	layer, err := tarball.LayerFromFile(input, tarball.WithMediaType(mediaType))
+	require.NoError(t, err)
+
+	return fakeLayer{layer}
+}
+
+type FakeDBOptions struct {
+	MediaType types.MediaType
+}
+
+func NewFakeDB(t *testing.T, dbPath string, opts FakeDBOptions) *oci.Artifact {
+	mediaType := lo.Ternary(opts.MediaType != "", opts.MediaType, defaultMediaType)
+	img := new(fakei.FakeImage)
+	img.LayersReturns([]v1.Layer{NewFakeLayer(t, dbPath, mediaType)}, nil)
+	img.ManifestReturns(&v1.Manifest{
+		Layers: []v1.Descriptor{
+			{
+				MediaType: mediaType,
+				Size:      100,
+				Digest: v1.Hash{
+					Algorithm: "sha256",
+					Hex:       "aec482bc254b5dd025d3eaf5bb35997d3dba783e394e8f91d5a415963151bfb8",
+				},
+				Annotations: map[string]string{
+					"org.opencontainers.image.title": "db.tar.gz",
+				},
+			},
+		},
+	}, nil)
+
+	// Mock OCI artifact
+	opt := ftypes.RegistryOptions{
+		Insecure: false,
+	}
+	art, err := oci.NewArtifact("dummy", true, opt, oci.WithImage(img))
+	require.NoError(t, err)
+
+	return art
+}
+
+func ArchiveDir(t *testing.T, dir string) string {
+	tmpDBPath := filepath.Join(t.TempDir(), "db.tar")
+	f, err := os.Create(tmpDBPath)
+	require.NoError(t, err)
+	defer f.Close()
+
+	tr := tar.NewWriter(f)
+	defer tr.Close()
+
+	err = tr.AddFS(os.DirFS(dir))
+	require.NoError(t, err)
+
+	return tmpDBPath
+}

pkg/commands/operation/operation.go

@@ -117,7 +117,7 @@ func DownloadDB(ctx context.Context, appVersion, cacheDir string, dbRepository n
 	defer mu.Unlock()
 
 	client := db.NewClient(cacheDir, quiet, db.WithDBRepository(dbRepository))
-	needsUpdate, err := client.NeedsUpdate(appVersion, skipUpdate)
+	needsUpdate, err := client.NeedsUpdate(ctx, appVersion, skipUpdate)
 	if err != nil {
 		return xerrors.Errorf("database error: %w", err)
 	}

pkg/db/db.go

@@ -9,10 +9,10 @@ import (
 	"github.com/google/go-containerregistry/pkg/name"
 	"github.com/google/go-containerregistry/pkg/v1/remote/transport"
 	"golang.org/x/xerrors"
-	"k8s.io/utils/clock"
 
 	"github.com/aquasecurity/trivy-db/pkg/db"
 	"github.com/aquasecurity/trivy-db/pkg/metadata"
+	"github.com/aquasecurity/trivy/pkg/clock"
 	"github.com/aquasecurity/trivy/pkg/fanal/types"
 	"github.com/aquasecurity/trivy/pkg/log"
 	"github.com/aquasecurity/trivy/pkg/oci"
@@ -28,15 +28,8 @@ var (
 	defaultRepository, _ = name.NewTag(DefaultRepository)
 )
 
-// Operation defines the DB operations
-type Operation interface {
-	NeedsUpdate(cliVersion string, skip bool) (need bool, err error)
-	Download(ctx context.Context, dst string, opt types.RegistryOptions) (err error)
-}
-
 type options struct {
 	artifact     *oci.Artifact
-	clock        clock.Clock
 	dbRepository name.Reference
 }
 
@@ -57,13 +50,6 @@ func WithDBRepository(dbRepository name.Reference) Option {
 	}
 }
 
-// WithClock takes a clock
-func WithClock(c clock.Clock) Option {
-	return func(opts *options) {
-		opts.clock = c
-	}
-}
-
 // Client implements DB operations
 type Client struct {
 	*options
@@ -76,7 +62,6 @@ type Client struct {
 // NewClient is the factory method for DB client
 func NewClient(cacheDir string, quiet bool, opts ...Option) *Client {
 	o := &options{
-		clock:        clock.RealClock{},
 		dbRepository: defaultRepository,
 	}
 
@@ -93,7 +78,7 @@ func NewClient(cacheDir string, quiet bool, opts ...Option) *Client {
 }
 
 // NeedsUpdate check is DB needs update
-func (c *Client) NeedsUpdate(cliVersion string, skip bool) (bool, error) {
+func (c *Client) NeedsUpdate(ctx context.Context, cliVersion string, skip bool) (bool, error) {
 	meta, err := c.metadata.Get()
 	if err != nil {
 		log.Debug("There is no valid metadata file", log.Err(err))
@@ -124,7 +109,7 @@ func (c *Client) NeedsUpdate(cliVersion string, skip bool) (bool, error) {
 		return true, nil
 	}
 
-	return !c.isNewDB(meta), nil
+	return !c.isNewDB(ctx, meta), nil
 }
 
 func (c *Client) validate(meta metadata.Metadata) error {
@@ -136,13 +121,14 @@ func (c *Client) validate(meta metadata.Metadata) error {
 	return nil
 }
 
-func (c *Client) isNewDB(meta metadata.Metadata) bool {
-	if c.clock.Now().Before(meta.NextUpdate) {
+func (c *Client) isNewDB(ctx context.Context, meta metadata.Metadata) bool {
+	now := clock.Now(ctx)
+	if now.Before(meta.NextUpdate) {
 		log.Debug("DB update was skipped because the local DB is the latest")
 		return true
 	}
 
-	if c.clock.Now().Before(meta.DownloadedAt.Add(time.Hour)) {
+	if now.Before(meta.DownloadedAt.Add(time.Hour)) {
 		log.Debug("DB update was skipped because the local DB was downloaded during the last hour")
 		return true
 	}
@@ -165,13 +151,13 @@ func (c *Client) Download(ctx context.Context, dst string, opt types.RegistryOpt
 		return xerrors.Errorf("database download error: %w", err)
 	}
 
-	if err = c.updateDownloadedAt(dst); err != nil {
+	if err = c.updateDownloadedAt(ctx, dst); err != nil {
 		return xerrors.Errorf("failed to update downloaded_at: %w", err)
 	}
 	return nil
 }
 
-func (c *Client) updateDownloadedAt(dst string) error {
+func (c *Client) updateDownloadedAt(ctx context.Context, dst string) error {
 	log.Debug("Updating database metadata...")
 
 	// We have to initialize a metadata client here
@@ -182,7 +168,7 @@ func (c *Client) updateDownloadedAt(dst string) error {
 		return xerrors.Errorf("unable to get metadata: %w", err)
 	}
 
-	meta.DownloadedAt = c.clock.Now().UTC()
+	meta.DownloadedAt = clock.Now(ctx).UTC()
 	if err = client.Update(meta); err != nil {
 		return xerrors.Errorf("failed to update metadata: %w", err)
 	}