Skip to content

Commit 7848141

Browse files
amisevskamisevsk
committed
Add tests for jmm models
Test that listing models works as expected and that bytes -> human readable format conversion produces expected results.
1 parent 10197aa commit 7848141

File tree

3 files changed

+212
-17
lines changed

3 files changed

+212
-17
lines changed

‎pkg/cmd/models/cmd.go

+6-1
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package models
22

33
import (
44
"fmt"
5+
"jmm/pkg/lib/storage"
56

67
"github.com/spf13/cobra"
78
"github.com/spf13/viper"
@@ -50,10 +51,14 @@ func RunCommand(options *ModelsOptions) func(*cobra.Command, []string) {
5051
fmt.Println(err)
5152
return
5253
}
53-
err = listModels(options)
54+
store := storage.NewLocalStore(opts.configHome)
55+
56+
summary, err := listModels(store)
5457
if err != nil {
5558
fmt.Println(err)
5659
return
5760
}
61+
62+
fmt.Println(summary)
5863
}
5964
}

‎pkg/cmd/models/models.go

+25-16
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,13 @@ Copyright © 2024 Jozu.com
44
package models
55

66
import (
7+
"bytes"
78
"context"
89
"encoding/json"
910
"fmt"
1011
"jmm/pkg/artifact"
1112
"jmm/pkg/lib/storage"
12-
"os"
13+
"math"
1314
"text/tabwriter"
1415

1516
"github.com/opencontainers/go-digest"
@@ -21,23 +22,23 @@ const (
2122
ModelsTableFmt = "%s\t%s\t%s\t%s\t"
2223
)
2324

24-
func listModels(opts *ModelsOptions) error {
25-
store := storage.NewLocalStore(opts.configHome)
25+
func listModels(store storage.Store) (string, error) {
2626
index, err := store.ParseIndexJson()
2727
if err != nil {
28-
return err
28+
return "", err
2929
}
3030

3131
manifests, err := manifestsFromIndex(index, store)
3232
if err != nil {
33-
return err
33+
return "", err
3434
}
3535

36-
if err := printManifestsSummary(manifests, store); err != nil {
37-
return err
36+
summary, err := printManifestsSummary(manifests, store)
37+
if err != nil {
38+
return "", err
3839
}
3940

40-
return nil
41+
return summary, nil
4142
}
4243

4344
func manifestsFromIndex(index *ocispec.Index, store storage.Store) (map[digest.Digest]ocispec.Manifest, error) {
@@ -68,19 +69,20 @@ func readManifestConfig(manifest *ocispec.Manifest, store storage.Store) (*artif
6869
return config, nil
6970
}
7071

71-
func printManifestsSummary(manifests map[digest.Digest]ocispec.Manifest, store storage.Store) error {
72-
tw := tabwriter.NewWriter(os.Stdout, 0, 2, 3, ' ', 0)
72+
func printManifestsSummary(manifests map[digest.Digest]ocispec.Manifest, store storage.Store) (string, error) {
73+
buf := bytes.Buffer{}
74+
tw := tabwriter.NewWriter(&buf, 0, 2, 3, ' ', 0)
7375
fmt.Fprintln(tw, ModelsTableHeader)
7476
for digest, manifest := range manifests {
7577
// TODO: filter this list for manifests we're interested in (build needs to set a manifest mediaType/artifactType)
7678
line, err := getManifestInfoLine(digest, &manifest, store)
7779
if err != nil {
78-
return err
80+
return "", err
7981
}
8082
fmt.Fprintln(tw, line)
8183
}
8284
tw.Flush()
83-
return nil
85+
return buf.String(), nil
8486
}
8587

8688
func getManifestInfoLine(digest digest.Digest, manifest *ocispec.Manifest, store storage.Store) (string, error) {
@@ -103,17 +105,24 @@ func formatBytes(i int64) string {
103105
return "0 B"
104106
}
105107

106-
suffixes := []string{"B", "KiB", "MiB", "GiB", "TiB", "PiB"}
108+
if i < 1024 {
109+
// Catch bytes to avoid printing fractional amounts of bytes e.g. 123.0 bytes
110+
return fmt.Sprintf("%d B", i)
111+
}
112+
113+
suffixes := []string{"KiB", "MiB", "GiB", "TiB"}
107114
unit := float64(1024)
108115

109-
size := float64(i)
116+
size := float64(i) / unit
110117
for _, suffix := range suffixes {
111118
if size < unit {
112-
return fmt.Sprintf("%.1f %s", size, suffix)
119+
// Round down to the nearest tenth of a unit to avoid e.g. 1MiB - 1B = 1024KiB
120+
niceSize := math.Floor(size*10) / 10
121+
return fmt.Sprintf("%.1f %s", niceSize, suffix)
113122
}
114123
size = size / unit
115124
}
116125

117-
// Fall back to printing 1000's of PiB
126+
// Fall back to printing whatever's left as PiB
118127
return fmt.Sprintf("%.1f PiB", size)
119128
}

‎pkg/cmd/models/models_test.go

+181
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,181 @@
1+
package models
2+
3+
import (
4+
"fmt"
5+
"jmm/pkg/artifact"
6+
"jmm/pkg/lib/constants"
7+
internal "jmm/pkg/lib/testing"
8+
"testing"
9+
10+
"github.com/opencontainers/go-digest"
11+
ocispec "github.com/opencontainers/image-spec/specs-go/v1"
12+
"github.com/stretchr/testify/assert"
13+
)
14+
15+
func TestListModels(t *testing.T) {
16+
tests := []struct {
17+
name string
18+
manifests map[digest.Digest]ocispec.Manifest
19+
configs map[digest.Digest]artifact.JozuFile
20+
index *ocispec.Index
21+
expectedOutputRegexps []string
22+
expectErrRegexp string
23+
}{
24+
{
25+
name: "Cannot read index.json",
26+
index: nil,
27+
expectErrRegexp: "artifact not found",
28+
},
29+
{
30+
name: "Cannot find manifest from index.json",
31+
index: &ocispec.Index{
32+
Manifests: []ocispec.Descriptor{
33+
ManifestDesc("manifestA", true),
34+
ManifestDesc("manifestNotFound", true),
35+
},
36+
},
37+
manifests: map[digest.Digest]ocispec.Manifest{
38+
"manifestA": Manifest("configA", "layerA"),
39+
"manifestB": Manifest("configB", "layerB"),
40+
},
41+
configs: map[digest.Digest]artifact.JozuFile{
42+
"configA": Config("maintainerA", "formatA"),
43+
"configB": Config("maintainerB", "formatB"),
44+
},
45+
expectErrRegexp: "failed to read manifest manifestNotFound.*",
46+
},
47+
{
48+
name: "Cannot find config in manifest",
49+
index: &ocispec.Index{
50+
Manifests: []ocispec.Descriptor{
51+
ManifestDesc("manifestA", true),
52+
ManifestDesc("manifestB", true),
53+
},
54+
},
55+
manifests: map[digest.Digest]ocispec.Manifest{
56+
"manifestA": Manifest("configA", "layerA"),
57+
"manifestB": Manifest("configNotFound", "layerB"),
58+
},
59+
configs: map[digest.Digest]artifact.JozuFile{
60+
"configA": Config("maintainerA", "formatA"),
61+
"configB": Config("maintainerB", "formatB"),
62+
},
63+
expectErrRegexp: "failed to read config.*",
64+
},
65+
{
66+
name: "Prints summary of for each manifest line (layers are 1024 bytes)",
67+
index: &ocispec.Index{
68+
Manifests: []ocispec.Descriptor{
69+
ManifestDesc("manifestA", true),
70+
ManifestDesc("manifestB", true),
71+
},
72+
},
73+
manifests: map[digest.Digest]ocispec.Manifest{
74+
"manifestA": Manifest("configA", "layerA"),
75+
"manifestB": Manifest("configB", "layerB1", "layerB2", "layerB3"),
76+
},
77+
configs: map[digest.Digest]artifact.JozuFile{
78+
"configA": Config("maintainerA", "formatA"),
79+
"configB": Config("maintainerB", "formatB"),
80+
},
81+
expectedOutputRegexps: []string{
82+
"manifestA.*maintainerA.*formatA.*1.0 KiB",
83+
"manifestB.*maintainerB.*formatB.*3.0 KiB",
84+
},
85+
},
86+
}
87+
for _, tt := range tests {
88+
t.Run(tt.name, func(t *testing.T) {
89+
testStore := &internal.TestStore{
90+
Manifests: tt.manifests,
91+
Configs: tt.configs,
92+
Index: tt.index,
93+
}
94+
summary, err := listModels(testStore)
95+
if tt.expectErrRegexp != "" {
96+
// Should be error
97+
assert.Empty(t, summary, "Should not output summary on error")
98+
if assert.Error(t, err, "Should return an error") {
99+
return
100+
}
101+
assert.Regexp(t, tt.expectErrRegexp, err.Error())
102+
} else {
103+
if !assert.NoError(t, err, "Should not return an error") {
104+
return
105+
}
106+
for _, line := range tt.expectedOutputRegexps {
107+
// Assert all lines in expected output are somewhere in the summary
108+
assert.Regexp(t, line, summary)
109+
}
110+
}
111+
})
112+
}
113+
}
114+
115+
func TestFormatBytes(t *testing.T) {
116+
tests := []struct {
117+
input int64
118+
output string
119+
}{
120+
{input: 0, output: "0 B"},
121+
{input: 500, output: "500 B"},
122+
{input: 1<<10 - 1, output: "1023 B"},
123+
{input: 1 << 10, output: "1.0 KiB"},
124+
{input: 4.5 * (1 << 10), output: "4.5 KiB"},
125+
{input: 1<<20 - 1, output: "1023.9 KiB"},
126+
{input: 1 << 20, output: "1.0 MiB"},
127+
{input: 6.5 * (1 << 20), output: "6.5 MiB"},
128+
{input: 1<<30 - 1, output: "1023.9 MiB"},
129+
{input: 1 << 30, output: "1.0 GiB"},
130+
{input: 1 << 40, output: "1.0 TiB"},
131+
{input: 1 << 50, output: "1.0 PiB"},
132+
{input: 500 * (1 << 50), output: "500.0 PiB"},
133+
{input: 1 << 60, output: "1024.0 PiB"},
134+
}
135+
for idx, tt := range tests {
136+
t.Run(fmt.Sprintf("test %d", idx), func(t *testing.T) {
137+
output := formatBytes(tt.input)
138+
assert.Equalf(t, tt.output, output, "Should convert %d to %s", tt.input, tt.output)
139+
})
140+
}
141+
}
142+
143+
func Manifest(configDigest string, layerDigests ...string) ocispec.Manifest {
144+
manifest := ocispec.Manifest{
145+
Config: ocispec.Descriptor{
146+
MediaType: constants.ModelConfigMediaType,
147+
Digest: digest.Digest(configDigest),
148+
},
149+
}
150+
for _, layerDigest := range layerDigests {
151+
manifest.Layers = append(manifest.Layers, ocispec.Descriptor{
152+
MediaType: constants.ModelLayerMediaType,
153+
Digest: digest.Digest(layerDigest),
154+
Size: 1024,
155+
})
156+
}
157+
158+
return manifest
159+
}
160+
161+
func Config(maintainer, format string) artifact.JozuFile {
162+
config := artifact.JozuFile{
163+
Maintainer: maintainer,
164+
ModelFormat: format,
165+
}
166+
167+
return config
168+
}
169+
170+
func ManifestDesc(digestStr string, valid bool) ocispec.Descriptor {
171+
size := internal.ValidSize
172+
if !valid {
173+
size = internal.InvalidSize
174+
}
175+
176+
return ocispec.Descriptor{
177+
Digest: digest.Digest(digestStr),
178+
MediaType: ocispec.MediaTypeImageManifest,
179+
Size: size,
180+
}
181+
}

0 commit comments

Comments
 (0)