feat: add unit tests (#15)

Signed-off-by: chlins <[email protected]>

Author: chlins

Diff for: .github/workflows/ci.yaml

@@ -10,7 +10,7 @@ on:
   schedule:
     - cron: '0 4 * * *'
 
-permissions:  
+permissions:
   contents: read
 
 jobs:
@@ -31,4 +31,5 @@ jobs:
 
       - name: Run Unit tests
         run: |-
-          echo running unit tests
+          go version
+          go test -v ./...

Diff for: .golangci.yml

@@ -1,6 +1,9 @@
 run:
   deadline: 3m
   modules-download-mode: readonly
+  skip-dirs:
+    - test/mocks
+
 
 linters-settings:
   gocyclo:

Diff for: .mockery.yaml

@@ -0,0 +1,21 @@
+with-expecter: true
+boilerplate-file: copyright.txt
+outpkg: "{{.PackageName}}"
+mockname: "{{.InterfaceName}}"
+filename: "{{.InterfaceName | snakecase}}.go"
+packages:
+  github.com/CloudNativeAI/modctl/pkg/backend:
+    interfaces:
+      Backend:
+        config:
+          dir: test/mocks/backend
+  github.com/CloudNativeAI/modctl/pkg/storage:
+    interfaces:
+      Storage:
+        config:
+          dir: test/mocks/storage
+  github.com/CloudNativeAI/modctl/pkg/modelfile:
+    interfaces:
+      Modelfile:
+        config:
+          dir: test/mocks/modelfile

Diff for: copyright.txt

@@ -0,0 +1,15 @@
+/*
+ *     Copyright 2024 The CNAI Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */

Diff for: go.mod

@@ -83,6 +83,7 @@ require (
 	github.com/spf13/afero v1.11.0 // indirect
 	github.com/spf13/cast v1.7.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
+	github.com/stretchr/objx v0.5.2 // indirect
 	github.com/subosito/gotenv v1.6.0 // indirect
 	github.com/xeipuuv/gojsonpointer v0.0.0-20190905194746-02993c407bfb // indirect
 	github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 // indirect

Diff for: pkg/backend/build_test.go

@@ -0,0 +1,66 @@
+/*
+ *     Copyright 2024 The CNAI Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package backend
+
+import (
+	"testing"
+	"time"
+
+	modelspec "github.com/CloudNativeAI/modctl/pkg/spec"
+	"github.com/CloudNativeAI/modctl/test/mocks/modelfile"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestManifestAnnotation(t *testing.T) {
+	modelfile := &modelfile.Modelfile{}
+	modelfile.On("GetArch").Return("test-arch")
+	modelfile.On("GetFamily").Return("test-family")
+	modelfile.On("GetName").Return("test-model")
+	modelfile.On("GetFormat").Return("test-format")
+	modelfile.On("GetParamsize").Return("12345")
+	modelfile.On("GetPrecision").Return("FP32")
+	modelfile.On("GetQuantization").Return("INT8")
+
+	annotations := manifestAnnotation(modelfile)
+
+	assert.Equal(t, "test-arch", annotations[modelspec.AnnotationArchitecture])
+	assert.Equal(t, "test-family", annotations[modelspec.AnnotationFamily])
+	assert.Equal(t, "test-model", annotations[modelspec.AnnotationName])
+	assert.Equal(t, "test-format", annotations[modelspec.AnnotationFormat])
+	assert.Equal(t, "12345", annotations[modelspec.AnnotationParamSize])
+	assert.Equal(t, "FP32", annotations[modelspec.AnnotationPrecision])
+	assert.Equal(t, "INT8", annotations[modelspec.AnnotationQuantization])
+
+	createdTime, err := time.Parse(time.RFC3339, annotations[modelspec.AnnotationCreated])
+	assert.NoError(t, err)
+	assert.WithinDuration(t, time.Now(), createdTime, time.Minute)
+}
+
+func TestGetProcessors(t *testing.T) {
+	modelfile := &modelfile.Modelfile{}
+	modelfile.On("GetConfigs").Return([]string{"config1", "config2"})
+	modelfile.On("GetModels").Return([]string{"model1", "model2"})
+
+	processors := getProcessors(modelfile)
+
+	assert.Len(t, processors, 4)
+	assert.Equal(t, "license", processors[0].Name())
+	assert.Equal(t, "readme", processors[1].Name())
+	assert.Equal(t, "model_config", processors[2].Name())
+	assert.Equal(t, "model", processors[3].Name())
+}

Diff for: pkg/backend/list_test.go

@@ -0,0 +1,65 @@
+/*
+ *     Copyright 2024 The CNAI Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package backend
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/CloudNativeAI/modctl/test/mocks/storage"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+)
+
+func TestList(t *testing.T) {
+	mockStore := &storage.Storage{}
+	b := &backend{store: mockStore}
+	ctx := context.Background()
+	repos := []string{"example.com/repo1", "example.com/repo2"}
+	tags := []string{"tag1", "tag2"}
+	createdAt := time.Now().Format(time.RFC3339)
+	manifest := ocispec.Manifest{
+		Layers: []ocispec.Descriptor{
+			{Size: 1024},
+			{Size: 1024},
+		},
+		Config: ocispec.Descriptor{Size: 1024},
+		Annotations: map[string]string{
+			"org.cnai.model.created": createdAt,
+		},
+	}
+	manifestRaw, err := json.Marshal(manifest)
+	assert.NoError(t, err)
+
+	mockStore.On("ListRepositories", ctx).Return(repos, nil)
+	mockStore.On("ListTags", ctx, repos[0]).Return(tags, nil)
+	mockStore.On("ListTags", ctx, repos[1]).Return(tags, nil)
+	mockStore.On("PullManifest", ctx, mock.Anything, mock.Anything).Return(manifestRaw, "sha256:1234567890abcdef", nil)
+
+	artifacts, err := b.List(ctx)
+	assert.NoError(t, err, "list failed")
+	assert.Len(t, artifacts, 4, "unexpected number of artifacts")
+	assert.Equal(t, repos[0], artifacts[0].Repository, "unexpected repository")
+	assert.Equal(t, tags[0], artifacts[0].Tag, "unexpected tag")
+	assert.Equal(t, "sha256:1234567890abcdef", artifacts[0].Digest, "unexpected digest")
+	assert.Equal(t, int64(3*1024+len(manifestRaw)), artifacts[0].Size, "unexpected size")
+	assert.Equal(t, createdAt, artifacts[0].CreatedAt.Format(time.RFC3339), "unexpected created at")
+}

Diff for: pkg/backend/processor/license.go

@@ -35,6 +35,10 @@ func NewLicenseProcessor() Processor {
 // licenseProcessor is the processor to process the LICENSE file.
 type licenseProcessor struct{}
 
+func (p *licenseProcessor) Name() string {
+	return "license"
+}
+
 func (p *licenseProcessor) Identify(_ context.Context, path string, info os.FileInfo) bool {
 	return info.Name() == "LICENSE" || info.Name() == "LICENSE.txt"
 }

Diff for: pkg/backend/processor/license_test.go

@@ -0,0 +1,76 @@
+/*
+ *     Copyright 2024 The CNAI Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package processor
+
+import (
+	"context"
+	"testing"
+	"testing/fstest"
+
+	modelspec "github.com/CloudNativeAI/modctl/pkg/spec"
+	"github.com/CloudNativeAI/modctl/test/mocks/storage"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+)
+
+func TestLicenseProcessor_Name(t *testing.T) {
+	p := NewLicenseProcessor()
+	assert.Equal(t, "license", p.Name())
+}
+
+func TestLicenseProcessor_Identify(t *testing.T) {
+	p := NewLicenseProcessor()
+	mockFS := fstest.MapFS{
+		"LICENSE":     &fstest.MapFile{},
+		"LICENSE.txt": &fstest.MapFile{},
+		"README.md":   &fstest.MapFile{},
+	}
+	info, err := mockFS.Stat("LICENSE")
+	assert.NoError(t, err)
+	assert.True(t, p.Identify(context.Background(), "LICENSE", info))
+
+	info, err = mockFS.Stat("LICENSE.txt")
+	assert.NoError(t, err)
+	assert.True(t, p.Identify(context.Background(), "LICENSE.txt", info))
+
+	info, err = mockFS.Stat("README.md")
+	assert.NoError(t, err)
+	assert.False(t, p.Identify(context.Background(), "README.md", info))
+}
+
+func TestLicenseProcessor_Process(t *testing.T) {
+	p := NewLicenseProcessor()
+	ctx := context.Background()
+	mockStore := &storage.Storage{}
+	repo := "test-repo"
+	path := "LICENSE"
+	mockFS := fstest.MapFS{
+		"LICENSE": &fstest.MapFile{},
+	}
+	info, err := mockFS.Stat("LICENSE")
+	assert.NoError(t, err)
+
+	mockStore.On("PushBlob", ctx, repo, mock.Anything).Return("sha256:1234567890abcdef", int64(1024), nil)
+
+	desc, err := p.Process(ctx, mockStore, repo, path, info)
+	assert.NoError(t, err)
+	assert.NotNil(t, desc)
+	assert.Equal(t, "sha256:1234567890abcdef", desc.Digest.String())
+	assert.Equal(t, int64(1024), desc.Size)
+	assert.Equal(t, "true", desc.Annotations[modelspec.AnnotationLicense])
+}

Diff for: pkg/backend/processor/model.go

@@ -41,6 +41,10 @@ type modelProcessor struct {
 	models []string
 }
 
+func (p *modelProcessor) Name() string {
+	return "model"
+}
+
 func (p *modelProcessor) Identify(_ context.Context, path string, info os.FileInfo) bool {
 	for _, model := range p.models {
 		if matched, _ := regexp.MatchString(model, info.Name()); matched {

Diff for: pkg/backend/processor/model_config.go

@@ -41,6 +41,10 @@ type modelConfigProcessor struct {
 	configs []string
 }
 
+func (p *modelConfigProcessor) Name() string {
+	return "model_config"
+}
+
 func (p *modelConfigProcessor) Identify(_ context.Context, path string, info os.FileInfo) bool {
 	for _, config := range p.configs {
 		if matched, _ := regexp.MatchString(config, info.Name()); matched {

Diff for: pkg/backend/processor/model_config_test.go

@@ -0,0 +1,76 @@
+/*
+ *     Copyright 2024 The CNAI Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package processor
+
+import (
+	"context"
+	"testing"
+	"testing/fstest"
+
+	modelspec "github.com/CloudNativeAI/modctl/pkg/spec"
+	"github.com/CloudNativeAI/modctl/test/mocks/storage"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+)
+
+func TestModelConfigProcessor_Name(t *testing.T) {
+	p := NewModelConfigProcessor([]string{"config*"})
+	assert.Equal(t, "model_config", p.Name())
+}
+
+func TestModelConfigProcessor_Identify(t *testing.T) {
+	p := NewModelConfigProcessor([]string{"config*"})
+	mockFS := fstest.MapFS{
+		"config-1": &fstest.MapFile{},
+		"config-2": &fstest.MapFile{},
+		"model":    &fstest.MapFile{},
+	}
+	info, err := mockFS.Stat("config-1")
+	assert.NoError(t, err)
+	assert.True(t, p.Identify(context.Background(), "config-1", info))
+
+	info, err = mockFS.Stat("config-2")
+	assert.NoError(t, err)
+	assert.True(t, p.Identify(context.Background(), "config-2", info))
+
+	info, err = mockFS.Stat("model")
+	assert.NoError(t, err)
+	assert.False(t, p.Identify(context.Background(), "model", info))
+}
+
+func TestModelConfigProcessor_Process(t *testing.T) {
+	p := NewModelConfigProcessor([]string{"config*"})
+	ctx := context.Background()
+	mockStore := &storage.Storage{}
+	repo := "test-repo"
+	path := "config"
+	mockFS := fstest.MapFS{
+		"config": &fstest.MapFile{},
+	}
+	info, err := mockFS.Stat("config")
+	assert.NoError(t, err)
+
+	mockStore.On("PushBlob", ctx, repo, mock.Anything).Return("sha256:1234567890abcdef", int64(1024), nil)
+
+	desc, err := p.Process(ctx, mockStore, repo, path, info)
+	assert.NoError(t, err)
+	assert.NotNil(t, desc)
+	assert.Equal(t, "sha256:1234567890abcdef", desc.Digest.String())
+	assert.Equal(t, int64(1024), desc.Size)
+	assert.Equal(t, "true", desc.Annotations[modelspec.AnnotationConfig])
+}