Skip to content

Commit

Permalink
refactor: delete db mock (#6940)
Browse files Browse the repository at this point in the history
Signed-off-by: knqyf263 <[email protected]>
  • Loading branch information
knqyf263 authored Jun 20, 2024
1 parent 983ac15 commit e493fc9
Show file tree
Hide file tree
Showing 28 changed files with 194 additions and 371 deletions.
3 changes: 2 additions & 1 deletion integration/integration_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down
File renamed without changes.
84 changes: 84 additions & 0 deletions internal/dbtest/fake.go
Original file line number Diff line number Diff line change
@@ -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
}
2 changes: 1 addition & 1 deletion pkg/commands/operation/operation.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down
34 changes: 10 additions & 24 deletions pkg/db/db.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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
}

Expand All @@ -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
Expand All @@ -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,
}

Expand All @@ -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))
Expand Down Expand Up @@ -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 {
Expand All @@ -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
}
Expand All @@ -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
Expand All @@ -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)
}
Expand Down
Loading

0 comments on commit e493fc9

Please sign in to comment.