Rework storage to an interface; add tests for jmm models

pkg/cmd/build/build.go

@@ -8,6 +8,7 @@ import (
 	"os"
 
 	"jmm/pkg/artifact"
+	"jmm/pkg/lib/storage"
 
 	v1 "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/spf13/cobra"
@@ -106,7 +107,7 @@ func (options *BuildOptions) RunBuild() error {
 	model.Layers = append(model.Layers, layer)
 	model.Config = jozufile
 
-	store := artifact.NewArtifactStore(options.JozuHome)
+	store := storage.NewLocalStore(options.JozuHome)
 	var manifest *v1.Manifest
 	manifest, err = store.SaveModel(model)
 	if err != nil {

pkg/cmd/models/cmd.go

@@ -2,6 +2,7 @@ package models
 
 import (
 	"fmt"
+	"jmm/pkg/lib/storage"
 
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
@@ -50,10 +51,14 @@ func RunCommand(options *ModelsOptions) func(*cobra.Command, []string) {
 			fmt.Println(err)
 			return
 		}
-		err = listModels(options)
+		store := storage.NewLocalStore(opts.configHome)
+
+		summary, err := listModels(store)
 		if err != nil {
 			fmt.Println(err)
 			return
 		}
+
+		fmt.Println(summary)
 	}
 }

pkg/cmd/models/models.go

@@ -4,12 +4,13 @@ Copyright © 2024 Jozu.com
 package models
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
-	"io"
 	"jmm/pkg/artifact"
-	"os"
+	"jmm/pkg/lib/storage"
+	"math"
 	"text/tabwriter"
 
 	"github.com/opencontainers/go-digest"
@@ -21,33 +22,29 @@ const (
 	ModelsTableFmt    = "%s\t%s\t%s\t%s\t"
 )
 
-func listModels(opts *ModelsOptions) error {
-	store := artifact.NewArtifactStore(opts.configHome)
+func listModels(store storage.Store) (string, error) {
 	index, err := store.ParseIndexJson()
 	if err != nil {
-		return err
+		return "", err
 	}
 
 	manifests, err := manifestsFromIndex(index, store)
 	if err != nil {
-		return err
+		return "", err
 	}
 
-	if err := printManifestsSummary(manifests, store); err != nil {
-		return err
+	summary, err := printManifestsSummary(manifests, store)
+	if err != nil {
+		return "", err
 	}
 
-	return nil
+	return summary, nil
 }
 
-func manifestsFromIndex(index *ocispec.Index, store *artifact.Store) (map[digest.Digest]ocispec.Manifest, error) {
+func manifestsFromIndex(index *ocispec.Index, store storage.Store) (map[digest.Digest]ocispec.Manifest, error) {
 	manifests := map[digest.Digest]ocispec.Manifest{}
 	for _, manifestDesc := range index.Manifests {
-		manifestReader, err := store.Storage.Fetch(context.Background(), manifestDesc)
-		if err != nil {
-			return nil, fmt.Errorf("failed to get manifest %s: %w", manifestDesc.Digest, err)
-		}
-		manifestBytes, err := io.ReadAll(manifestReader)
+		manifestBytes, err := store.Fetch(context.Background(), manifestDesc)
 		if err != nil {
 			return nil, fmt.Errorf("failed to read manifest %s: %w", manifestDesc.Digest, err)
 		}
@@ -60,12 +57,8 @@ func manifestsFromIndex(index *ocispec.Index, store *artifact.Store) (map[digest
 	return manifests, nil
 }
 
-func readManifestConfig(manifest *ocispec.Manifest, store *artifact.Store) (*artifact.JozuFile, error) {
-	configReader, err := store.Storage.Fetch(context.Background(), manifest.Config)
-	if err != nil {
-		return nil, fmt.Errorf("failed to get config: %w", err)
-	}
-	configBytes, err := io.ReadAll(configReader)
+func readManifestConfig(manifest *ocispec.Manifest, store storage.Store) (*artifact.JozuFile, error) {
+	configBytes, err := store.Fetch(context.Background(), manifest.Config)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read config: %w", err)
 	}
@@ -76,22 +69,23 @@ func readManifestConfig(manifest *ocispec.Manifest, store *artifact.Store) (*art
 	return config, nil
 }
 
-func printManifestsSummary(manifests map[digest.Digest]ocispec.Manifest, store *artifact.Store) error {
-	tw := tabwriter.NewWriter(os.Stdout, 0, 2, 3, ' ', 0)
+func printManifestsSummary(manifests map[digest.Digest]ocispec.Manifest, store storage.Store) (string, error) {
+	buf := bytes.Buffer{}
+	tw := tabwriter.NewWriter(&buf, 0, 2, 3, ' ', 0)
 	fmt.Fprintln(tw, ModelsTableHeader)
 	for digest, manifest := range manifests {
 		// TODO: filter this list for manifests we're interested in (build needs to set a manifest mediaType/artifactType)
 		line, err := getManifestInfoLine(digest, &manifest, store)
 		if err != nil {
-			return err
+			return "", err
 		}
 		fmt.Fprintln(tw, line)
 	}
 	tw.Flush()
-	return nil
+	return buf.String(), nil
 }
 
-func getManifestInfoLine(digest digest.Digest, manifest *ocispec.Manifest, store *artifact.Store) (string, error) {
+func getManifestInfoLine(digest digest.Digest, manifest *ocispec.Manifest, store storage.Store) (string, error) {
 	config, err := readManifestConfig(manifest, store)
 	if err != nil {
 		return "", err
@@ -111,17 +105,24 @@ func formatBytes(i int64) string {
 		return "0 B"
 	}
 
-	suffixes := []string{"B", "KiB", "MiB", "GiB", "TiB", "PiB"}
+	if i < 1024 {
+		// Catch bytes to avoid printing fractional amounts of bytes e.g. 123.0 bytes
+		return fmt.Sprintf("%d B", i)
+	}
+
+	suffixes := []string{"KiB", "MiB", "GiB", "TiB"}
 	unit := float64(1024)
 
-	size := float64(i)
+	size := float64(i) / unit
 	for _, suffix := range suffixes {
 		if size < unit {
-			return fmt.Sprintf("%.1f %s", size, suffix)
+			// Round down to the nearest tenth of a unit to avoid e.g. 1MiB - 1B = 1024KiB
+			niceSize := math.Floor(size*10) / 10
+			return fmt.Sprintf("%.1f %s", niceSize, suffix)
 		}
 		size = size / unit
 	}
 
-	// Fall back to printing 1000's of PiB
+	// Fall back to printing whatever's left as PiB
 	return fmt.Sprintf("%.1f PiB", size)
 }

pkg/cmd/models/models_test.go

@@ -0,0 +1,181 @@
+package models
+
+import (
+	"fmt"
+	"jmm/pkg/artifact"
+	"jmm/pkg/lib/constants"
+	internal "jmm/pkg/lib/testing"
+	"testing"
+
+	"github.com/opencontainers/go-digest"
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestListModels(t *testing.T) {
+	tests := []struct {
+		name                  string
+		manifests             map[digest.Digest]ocispec.Manifest
+		configs               map[digest.Digest]artifact.JozuFile
+		index                 *ocispec.Index
+		expectedOutputRegexps []string
+		expectErrRegexp       string
+	}{
+		{
+			name:            "Cannot read index.json",
+			index:           nil,
+			expectErrRegexp: "artifact not found",
+		},
+		{
+			name: "Cannot find manifest from index.json",
+			index: &ocispec.Index{
+				Manifests: []ocispec.Descriptor{
+					ManifestDesc("manifestA", true),
+					ManifestDesc("manifestNotFound", true),
+				},
+			},
+			manifests: map[digest.Digest]ocispec.Manifest{
+				"manifestA": Manifest("configA", "layerA"),
+				"manifestB": Manifest("configB", "layerB"),
+			},
+			configs: map[digest.Digest]artifact.JozuFile{
+				"configA": Config("maintainerA", "formatA"),
+				"configB": Config("maintainerB", "formatB"),
+			},
+			expectErrRegexp: "failed to read manifest manifestNotFound.*",
+		},
+		{
+			name: "Cannot find config in manifest",
+			index: &ocispec.Index{
+				Manifests: []ocispec.Descriptor{
+					ManifestDesc("manifestA", true),
+					ManifestDesc("manifestB", true),
+				},
+			},
+			manifests: map[digest.Digest]ocispec.Manifest{
+				"manifestA": Manifest("configA", "layerA"),
+				"manifestB": Manifest("configNotFound", "layerB"),
+			},
+			configs: map[digest.Digest]artifact.JozuFile{
+				"configA": Config("maintainerA", "formatA"),
+				"configB": Config("maintainerB", "formatB"),
+			},
+			expectErrRegexp: "failed to read config.*",
+		},
+		{
+			name: "Prints summary of for each manifest line (layers are 1024 bytes)",
+			index: &ocispec.Index{
+				Manifests: []ocispec.Descriptor{
+					ManifestDesc("manifestA", true),
+					ManifestDesc("manifestB", true),
+				},
+			},
+			manifests: map[digest.Digest]ocispec.Manifest{
+				"manifestA": Manifest("configA", "layerA"),
+				"manifestB": Manifest("configB", "layerB1", "layerB2", "layerB3"),
+			},
+			configs: map[digest.Digest]artifact.JozuFile{
+				"configA": Config("maintainerA", "formatA"),
+				"configB": Config("maintainerB", "formatB"),
+			},
+			expectedOutputRegexps: []string{
+				"manifestA.*maintainerA.*formatA.*1.0 KiB",
+				"manifestB.*maintainerB.*formatB.*3.0 KiB",
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			testStore := &internal.TestStore{
+				Manifests: tt.manifests,
+				Configs:   tt.configs,
+				Index:     tt.index,
+			}
+			summary, err := listModels(testStore)
+			if tt.expectErrRegexp != "" {
+				// Should be error
+				assert.Empty(t, summary, "Should not output summary on error")
+				if assert.Error(t, err, "Should return an error") {
+					return
+				}
+				assert.Regexp(t, tt.expectErrRegexp, err.Error())
+			} else {
+				if !assert.NoError(t, err, "Should not return an error") {
+					return
+				}
+				for _, line := range tt.expectedOutputRegexps {
+					// Assert all lines in expected output are somewhere in the summary
+					assert.Regexp(t, line, summary)
+				}
+			}
+		})
+	}
+}
+
+func TestFormatBytes(t *testing.T) {
+	tests := []struct {
+		input  int64
+		output string
+	}{
+		{input: 0, output: "0 B"},
+		{input: 500, output: "500 B"},
+		{input: 1<<10 - 1, output: "1023 B"},
+		{input: 1 << 10, output: "1.0 KiB"},
+		{input: 4.5 * (1 << 10), output: "4.5 KiB"},
+		{input: 1<<20 - 1, output: "1023.9 KiB"},
+		{input: 1 << 20, output: "1.0 MiB"},
+		{input: 6.5 * (1 << 20), output: "6.5 MiB"},
+		{input: 1<<30 - 1, output: "1023.9 MiB"},
+		{input: 1 << 30, output: "1.0 GiB"},
+		{input: 1 << 40, output: "1.0 TiB"},
+		{input: 1 << 50, output: "1.0 PiB"},
+		{input: 500 * (1 << 50), output: "500.0 PiB"},
+		{input: 1 << 60, output: "1024.0 PiB"},
+	}
+	for idx, tt := range tests {
+		t.Run(fmt.Sprintf("test %d", idx), func(t *testing.T) {
+			output := formatBytes(tt.input)
+			assert.Equalf(t, tt.output, output, "Should convert %d to %s", tt.input, tt.output)
+		})
+	}
+}
+
+func Manifest(configDigest string, layerDigests ...string) ocispec.Manifest {
+	manifest := ocispec.Manifest{
+		Config: ocispec.Descriptor{
+			MediaType: constants.ModelConfigMediaType,
+			Digest:    digest.Digest(configDigest),
+		},
+	}
+	for _, layerDigest := range layerDigests {
+		manifest.Layers = append(manifest.Layers, ocispec.Descriptor{
+			MediaType: constants.ModelLayerMediaType,
+			Digest:    digest.Digest(layerDigest),
+			Size:      1024,
+		})
+	}
+
+	return manifest
+}
+
+func Config(maintainer, format string) artifact.JozuFile {
+	config := artifact.JozuFile{
+		Maintainer:  maintainer,
+		ModelFormat: format,
+	}
+
+	return config
+}
+
+func ManifestDesc(digestStr string, valid bool) ocispec.Descriptor {
+	size := internal.ValidSize
+	if !valid {
+		size = internal.InvalidSize
+	}
+
+	return ocispec.Descriptor{
+		Digest:    digest.Digest(digestStr),
+		MediaType: ocispec.MediaTypeImageManifest,
+		Size:      size,
+	}
+}

pkg/artifact/consts.go → pkg/lib/constants/consts.go

@@ -1,4 +1,4 @@
-package artifact
+package constants
 
 const (
 	// Media type for the model layer