From 3bc6cbe2abf5bae90c0c54f54d7a7fa30310e9d3 Mon Sep 17 00:00:00 2001 From: Angel Misevski Date: Tue, 14 Jan 2025 19:19:04 -0500 Subject: [PATCH 1/3] Add tests for generating Kitfiles. --- testing/generate-kitfile_test.go | 84 +++++++++++++++++++ .../test_directory-handling.yaml | 49 +++++++++++ .../test_empty-directory.yaml | 10 +++ .../test_existing-Kitfile.yaml | 17 ++++ .../test_metadata-files-dataset.yaml | 23 +++++ .../test_metadata-files-model.yaml | 33 ++++++++ .../test_model-and-adaptor.yaml | 21 +++++ .../test_multi-part-models.yaml | 22 +++++ .../kitfile-generation/test_type-detect.yaml | 46 ++++++++++ 9 files changed, 305 insertions(+) create mode 100644 testing/generate-kitfile_test.go create mode 100644 testing/testdata/kitfile-generation/test_directory-handling.yaml create mode 100644 testing/testdata/kitfile-generation/test_empty-directory.yaml create mode 100644 testing/testdata/kitfile-generation/test_existing-Kitfile.yaml create mode 100644 testing/testdata/kitfile-generation/test_metadata-files-dataset.yaml create mode 100644 testing/testdata/kitfile-generation/test_metadata-files-model.yaml create mode 100644 testing/testdata/kitfile-generation/test_model-and-adaptor.yaml create mode 100644 testing/testdata/kitfile-generation/test_multi-part-models.yaml create mode 100644 testing/testdata/kitfile-generation/test_type-detect.yaml diff --git a/testing/generate-kitfile_test.go b/testing/generate-kitfile_test.go new file mode 100644 index 000000000..cba1635d4 --- /dev/null +++ b/testing/generate-kitfile_test.go @@ -0,0 +1,84 @@ +// Copyright 2025 The KitOps 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. +// +// SPDX-License-Identifier: Apache-2.0 + +package testing + +import ( + "fmt" + "kitops/pkg/artifact" + "kitops/pkg/lib/constants" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" +) + +type kitfileGenTestCase struct { + // Set to filename when loading test + Name string + ModelName string `yaml:"modelName"` + Description string `yaml:"description"` + Files []string `yaml:"files"` + ExpectedKitfile *artifact.KitFile `yaml:"expectedKitfile"` +} + +func (t kitfileGenTestCase) withName(name string) kitfileGenTestCase { + t.Name = name + return t +} + +func TestKitfileGeneration(t *testing.T) { + cleanup := testPreflight(t) + defer cleanup(t) + tests := loadAllTestCasesOrPanic[kitfileGenTestCase](t, filepath.Join("testdata", "kitfile-generation")) + for _, tt := range tests { + t.Run(fmt.Sprintf("%s (%s)", tt.Name, tt.Description), func(t *testing.T) { + // Set up temporary directory for work + tmpDir, removeTmp := setupTempDir(t) + defer removeTmp() + + // Set up directory for KITOPS_HOME + contextPath := filepath.Join(tmpDir, ".kitops") + if err := os.MkdirAll(contextPath, 0755); err != nil { + t.Fatal(err) + } + t.Setenv(constants.KitopsHomeEnvVar, contextPath) + + // Set up separate directory for files + testPath := filepath.Join(tmpDir, "testdata") + if err := os.MkdirAll(testPath, 0755); err != nil { + t.Fatal(err) + } + + // Create files listed in test case + setupFiles(t, testPath, tt.Files) + + runCommand(t, expectNoError, "init", testPath, "--name", tt.ModelName, "--desc", tt.Description, "--force") + + actualKitfile := &artifact.KitFile{} + modelfile, err := os.Open(filepath.Join(testPath, constants.DefaultKitfileName)) + if !assert.NoError(t, err) { + return + } + defer modelfile.Close() + if err := actualKitfile.LoadModel(modelfile); !assert.NoError(t, err) { + return + } + assert.Equal(t, tt.ExpectedKitfile, actualKitfile, "Generated Kitfile should match") + }) + } +} diff --git a/testing/testdata/kitfile-generation/test_directory-handling.yaml b/testing/testdata/kitfile-generation/test_directory-handling.yaml new file mode 100644 index 000000000..ed54e0773 --- /dev/null +++ b/testing/testdata/kitfile-generation/test_directory-handling.yaml @@ -0,0 +1,49 @@ +description: "Handles directories" + +files: + # Special-cased names + - src/test-file.sh + - pkg/test-file.sh + - lib/test-file.sh + - build/test-file.sh + - docs/documentation-file + # directories with only models + - model-contents/my-model1.gguf + - model-contents/my-model2.gguf + - model-contents/z-metadata.json + # directories with only datasets + - dataset-contents/dataset-001.zip + - dataset-contents/dataset-002.zip + - dataset-contents/dataset-003.zip + - dataset-contents/z-metadata.json + # documentation-only directory with weird name + - my-files/doc1.md + - my-files/doc2.md + - my-files/new.pdf + # mixed contents to be handled as code + - mixed-contents/documentation.md + - mixed-contents/dataset.csv + - mixed-contents/meta.json + +modelName: test-directory-handling +expectedKitfile: + manifestVersion: "1.0.0" + package: + name: test-directory-handling + description: "Handles directories" + model: + path: model-contents/my-model1.gguf + parts: + - path: model-contents/my-model2.gguf + - path: model-contents/z-metadata.json + datasets: + - path: dataset-contents + docs: + - path: docs + - path: my-files + code: + - path: build + - path: lib + - path: pkg + - path: src + - path: mixed-contents diff --git a/testing/testdata/kitfile-generation/test_empty-directory.yaml b/testing/testdata/kitfile-generation/test_empty-directory.yaml new file mode 100644 index 000000000..828eefe0e --- /dev/null +++ b/testing/testdata/kitfile-generation/test_empty-directory.yaml @@ -0,0 +1,10 @@ +description: "Generates empty Kitfile in empty directory" + +files: [] + +modelName: test-empty +expectedKitfile: + manifestVersion: "1.0.0" + package: + name: test-empty + description: "Generates empty Kitfile in empty directory" diff --git a/testing/testdata/kitfile-generation/test_existing-Kitfile.yaml b/testing/testdata/kitfile-generation/test_existing-Kitfile.yaml new file mode 100644 index 000000000..0f32727fa --- /dev/null +++ b/testing/testdata/kitfile-generation/test_existing-Kitfile.yaml @@ -0,0 +1,17 @@ +description: "Handles existing Kitfiles" + +files: + - test-model.npy + - Kitfile + - kitfile + - .kitfile + +modelName: test-existing-kitfile +expectedKitfile: + manifestVersion: "1.0.0" + package: + name: test-existing-kitfile + description: "Handles existing Kitfiles" + model: + name: test-model + path: test-model.npy diff --git a/testing/testdata/kitfile-generation/test_metadata-files-dataset.yaml b/testing/testdata/kitfile-generation/test_metadata-files-dataset.yaml new file mode 100644 index 000000000..f2fe0fd59 --- /dev/null +++ b/testing/testdata/kitfile-generation/test_metadata-files-dataset.yaml @@ -0,0 +1,23 @@ +description: "Attaches metadata files to datasets when no model" + +files: + - a-dataset.tar + - b-dataset.csv + - c-dataset.zip + - x-metadata.json + - y-metadata.yaml + - z-metadata.xml + +modelName: test-dataset-metadata +expectedKitfile: + manifestVersion: "1.0.0" + package: + name: test-dataset-metadata + description: "Attaches metadata files to datasets when no model" + datasets: + - path: a-dataset.tar + - path: b-dataset.csv + - path: c-dataset.zip + - path: x-metadata.json + - path: y-metadata.yaml + - path: z-metadata.xml diff --git a/testing/testdata/kitfile-generation/test_metadata-files-model.yaml b/testing/testdata/kitfile-generation/test_metadata-files-model.yaml new file mode 100644 index 000000000..ba0dba227 --- /dev/null +++ b/testing/testdata/kitfile-generation/test_metadata-files-model.yaml @@ -0,0 +1,33 @@ +description: "Attaches metadata files to Model when present" + +files: + - a-model.gguf + - b-model.safetensors + - c-model.pkl + - d-model.pb + - a-dataset.tar + - b-dataset.csv + - c-dataset.zip + - x-metadata.json + - y-metadata.yaml + - z-metadata.xml + +modelName: test-model-metadata +expectedKitfile: + manifestVersion: "1.0.0" + package: + name: test-model-metadata + description: "Attaches metadata files to Model when present" + model: + path: a-model.gguf + parts: + - path: b-model.safetensors + - path: c-model.pkl + - path: d-model.pb + - path: x-metadata.json + - path: y-metadata.yaml + - path: z-metadata.xml + datasets: + - path: a-dataset.tar + - path: b-dataset.csv + - path: c-dataset.zip diff --git a/testing/testdata/kitfile-generation/test_model-and-adaptor.yaml b/testing/testdata/kitfile-generation/test_model-and-adaptor.yaml new file mode 100644 index 000000000..43bea79c7 --- /dev/null +++ b/testing/testdata/kitfile-generation/test_model-and-adaptor.yaml @@ -0,0 +1,21 @@ +description: "Handles models with an adaptor" + +files: + - model-adaptor1.gguf + - model-adaptor2.gguf + # Heuristically, this needs to be significantly bigger than the other files + # Since we store the filename in the file, we'll just make the filename long + - zzzzzzz-main-model-larger-than-normal-files.gguf + +modelName: test-model-and-adaptor +expectedKitfile: + manifestVersion: "1.0.0" + package: + name: test-model-and-adaptor + description: "Handles models with an adaptor" + model: + name: zzzzzzz-main-model-larger-than-normal-files + path: zzzzzzz-main-model-larger-than-normal-files.gguf + parts: + - path: model-adaptor1.gguf + - path: model-adaptor2.gguf diff --git a/testing/testdata/kitfile-generation/test_multi-part-models.yaml b/testing/testdata/kitfile-generation/test_multi-part-models.yaml new file mode 100644 index 000000000..b52be4ec2 --- /dev/null +++ b/testing/testdata/kitfile-generation/test_multi-part-models.yaml @@ -0,0 +1,22 @@ +description: "Handles multi-part models" + +files: + - model-0005.safetensors + - model-0004.safetensors + - model-0003.safetensors + - model-0002.safetensors + - model-0001.safetensors + +modelName: test-multipart +expectedKitfile: + manifestVersion: "1.0.0" + package: + name: test-multipart + description: "Handles multi-part models" + model: + path: model-0001.safetensors + parts: + - path: model-0002.safetensors + - path: model-0003.safetensors + - path: model-0004.safetensors + - path: model-0005.safetensors diff --git a/testing/testdata/kitfile-generation/test_type-detect.yaml b/testing/testdata/kitfile-generation/test_type-detect.yaml new file mode 100644 index 000000000..8a9a9f1e2 --- /dev/null +++ b/testing/testdata/kitfile-generation/test_type-detect.yaml @@ -0,0 +1,46 @@ +description: "Detects type of files based on extension" + +files: + - a-model.gguf + - b-model.safetensors + - c-model.pkl + - d-model.pb + - a-dataset.tar + - b-dataset.csv + - c-dataset.zip + - a-docs.md + - b-docs.adoc + - c-docs.pdf + # Specially-handled "documentation" files -- license and README + - README.md + - LICENSE + # Generic code files that should be caught in a catch-all section + - test.sh + - bootstrap.sh + +modelName: test-model-detect +expectedKitfile: + manifestVersion: "1.0.0" + package: + name: test-model-detect + description: "Detects type of files based on extension" + model: + path: a-model.gguf + parts: + - path: b-model.safetensors + - path: c-model.pkl + - path: d-model.pb + datasets: + - path: a-dataset.tar + - path: b-dataset.csv + - path: c-dataset.zip + docs: + - path: LICENSE + description: License file + - path: README.md + description: Readme file + - path: a-docs.md + - path: b-docs.adoc + - path: c-docs.pdf + code: + - path: "." From 4738e0941b3eb24dac04c97c4744aeb2494ee80c Mon Sep 17 00:00:00 2001 From: Angel Misevski Date: Thu, 16 Jan 2025 14:36:12 -0500 Subject: [PATCH 2/3] Use t.Cleanup in tests instead of defer --- testing/generate-kitfile_test.go | 13 ++++++++----- testing/modelkit-refs_test.go | 7 +++---- testing/pack-unpack_test.go | 10 ++++------ testing/remove_test.go | 21 +++++++-------------- testing/util_test.go | 18 ++++++++++-------- 5 files changed, 32 insertions(+), 37 deletions(-) diff --git a/testing/generate-kitfile_test.go b/testing/generate-kitfile_test.go index cba1635d4..418b7f3c0 100644 --- a/testing/generate-kitfile_test.go +++ b/testing/generate-kitfile_test.go @@ -42,14 +42,13 @@ func (t kitfileGenTestCase) withName(name string) kitfileGenTestCase { } func TestKitfileGeneration(t *testing.T) { - cleanup := testPreflight(t) - defer cleanup(t) + testPreflight(t) + tests := loadAllTestCasesOrPanic[kitfileGenTestCase](t, filepath.Join("testdata", "kitfile-generation")) for _, tt := range tests { t.Run(fmt.Sprintf("%s (%s)", tt.Name, tt.Description), func(t *testing.T) { // Set up temporary directory for work - tmpDir, removeTmp := setupTempDir(t) - defer removeTmp() + tmpDir := setupTempDir(t) // Set up directory for KITOPS_HOME contextPath := filepath.Join(tmpDir, ".kitops") @@ -74,7 +73,11 @@ func TestKitfileGeneration(t *testing.T) { if !assert.NoError(t, err) { return } - defer modelfile.Close() + t.Cleanup(func() { + if err := modelfile.Close(); err != nil { + t.Fatalf("Error closing model file: %s", err) + } + }) if err := actualKitfile.LoadModel(modelfile); !assert.NoError(t, err) { return } diff --git a/testing/modelkit-refs_test.go b/testing/modelkit-refs_test.go index 89ea9661d..cf3488da5 100644 --- a/testing/modelkit-refs_test.go +++ b/testing/modelkit-refs_test.go @@ -44,14 +44,13 @@ func (t modelkitRefTestcase) withName(name string) modelkitRefTestcase { } func TestModelKitReferences(t *testing.T) { - cleanup := testPreflight(t) - defer cleanup(t) + testPreflight(t) + tests := loadAllTestCasesOrPanic[modelkitRefTestcase](t, filepath.Join("testdata", "modelkit-refs")) for _, tt := range tests { t.Run(fmt.Sprintf("%s (%s)", tt.Name, tt.Description), func(t *testing.T) { // Set up temporary directory for work - tmpDir, removeTmp := setupTempDir(t) - defer removeTmp() + tmpDir := setupTempDir(t) // Set up directory for KITOPS_HOME contextPath := filepath.Join(tmpDir, ".kitops") diff --git a/testing/pack-unpack_test.go b/testing/pack-unpack_test.go index 415190b9d..fd3b51409 100644 --- a/testing/pack-unpack_test.go +++ b/testing/pack-unpack_test.go @@ -47,14 +47,13 @@ func (t packUnpackTestcase) withName(name string) packUnpackTestcase { // We work in a new temporary directory for each test to avoid interaction between // tests. func TestPackUnpack(t *testing.T) { - cleanup := testPreflight(t) - defer cleanup(t) + testPreflight(t) + tests := loadAllTestCasesOrPanic[packUnpackTestcase](t, filepath.Join("testdata", "pack-unpack")) for _, tt := range tests { t.Run(fmt.Sprintf("%s (%s)", tt.Name, tt.Description), func(t *testing.T) { // Set up temporary directory for work - tmpDir, removeTmp := setupTempDir(t) - defer removeTmp() + tmpDir := setupTempDir(t) // Set up paths to use for test modelKitPath, unpackPath, contextPath := setupTestDirs(t, tmpDir) @@ -76,8 +75,7 @@ func TestPackUnpack(t *testing.T) { } func TestPackReproducibility(t *testing.T) { - tmpDir, removeTmp := setupTempDir(t) - defer removeTmp() + tmpDir := setupTempDir(t) modelKitPath, _, contextPath := setupTestDirs(t, tmpDir) t.Setenv(constants.KitopsHomeEnvVar, contextPath) diff --git a/testing/remove_test.go b/testing/remove_test.go index a49771a7b..8f30fc9d2 100644 --- a/testing/remove_test.go +++ b/testing/remove_test.go @@ -39,8 +39,7 @@ model: func TestRemoveSingleModelkitTag(t *testing.T) { // Set up temporary directory for work - tmpDir, removeTmp := setupTempDir(t) - defer removeTmp() + tmpDir := setupTempDir(t) modelKitPath, _, contextPath := setupTestDirs(t, tmpDir) t.Setenv(constants.KitopsHomeEnvVar, contextPath) @@ -67,8 +66,7 @@ func TestRemoveSingleModelkitTag(t *testing.T) { func TestRemoveSingleModelkitDigest(t *testing.T) { // Set up temporary directory for work - tmpDir, removeTmp := setupTempDir(t) - defer removeTmp() + tmpDir := setupTempDir(t) modelKitPath, _, contextPath := setupTestDirs(t, tmpDir) t.Setenv(constants.KitopsHomeEnvVar, contextPath) @@ -96,8 +94,7 @@ func TestRemoveSingleModelkitDigest(t *testing.T) { func TestRemoveSingleModelkitNoTag(t *testing.T) { // Set up temporary directory for work - tmpDir, removeTmp := setupTempDir(t) - defer removeTmp() + tmpDir := setupTempDir(t) modelKitPath, _, contextPath := setupTestDirs(t, tmpDir) t.Setenv(constants.KitopsHomeEnvVar, contextPath) @@ -124,8 +121,7 @@ func TestRemoveSingleModelkitNoTag(t *testing.T) { func TestRemoveModelkitUntagsWhenMultiple(t *testing.T) { // Set up temporary directory for work - tmpDir, removeTmp := setupTempDir(t) - defer removeTmp() + tmpDir := setupTempDir(t) modelKitPath, _, contextPath := setupTestDirs(t, tmpDir) t.Setenv(constants.KitopsHomeEnvVar, contextPath) @@ -157,8 +153,7 @@ func TestRemoveModelkitUntagsWhenMultiple(t *testing.T) { func TestRemoveModelkitUntagsAllWhenDigest(t *testing.T) { // Set up temporary directory for work - tmpDir, removeTmp := setupTempDir(t) - defer removeTmp() + tmpDir := setupTempDir(t) modelKitPath, _, contextPath := setupTestDirs(t, tmpDir) t.Setenv(constants.KitopsHomeEnvVar, contextPath) @@ -191,8 +186,7 @@ func TestRemoveModelkitUntagsAllWhenDigest(t *testing.T) { func TestRemoveModelkitUntagged(t *testing.T) { // Set up temporary directory for work - tmpDir, removeTmp := setupTempDir(t) - defer removeTmp() + tmpDir := setupTempDir(t) modelKitPath, _, contextPath := setupTestDirs(t, tmpDir) t.Setenv(constants.KitopsHomeEnvVar, contextPath) @@ -234,8 +228,7 @@ func TestRemoveModelkitUntagged(t *testing.T) { func TestRemoveModelkitAll(t *testing.T) { // Set up temporary directory for work - tmpDir, removeTmp := setupTempDir(t) - defer removeTmp() + tmpDir := setupTempDir(t) modelKitPath, _, contextPath := setupTestDirs(t, tmpDir) t.Setenv(constants.KitopsHomeEnvVar, contextPath) diff --git a/testing/util_test.go b/testing/util_test.go index 9f4c0bff1..c722b9eb6 100644 --- a/testing/util_test.go +++ b/testing/util_test.go @@ -44,16 +44,16 @@ const ( // testPreflight should be called at the start of every test; it returns a function that // restores state (e.g. working directory) that may have been changed by executing commands. -func testPreflight(t *testing.T) func(t *testing.T) { +func testPreflight(t *testing.T) { wd, err := os.Getwd() if err != nil { t.Fatal(err) } - return func(t *testing.T) { + t.Cleanup(func() { if err := os.Chdir(wd); err != nil { t.Fatal(err) } - } + }) } // runCommand executes kit , saving stdout/stderr output to a buffer @@ -86,19 +86,21 @@ func runCommand(t *testing.T, e shouldExpectError, args ...string) string { return string(outlog) } -func setupTempDir(t *testing.T) (tmpDir string, removeTmpDir func()) { +// setupTempDir creates a temporary directory and returns its path. Directory is automatically +// cleaned up on test completion. +func setupTempDir(t *testing.T) (tmpDir string) { // Set up temporary directory for work tmpDir, err := os.MkdirTemp("", "kitops-testing-*") if !assert.NoError(t, err) { t.Fatalf("Could not create temporary directory: %s", err) } - removeTmpDir = func() { + t.Logf("Using temp directory: %s", tmpDir) + t.Cleanup(func() { if err := os.RemoveAll(tmpDir); err != nil { t.Logf("Error removing temp dir: %s", err) } - } - t.Logf("Using temp directory: %s", tmpDir) - return tmpDir, removeTmpDir + }) + return tmpDir } // setupTestDirs generates the test directories used for storing $KIT_HOME, the original modelkit From e998dc24349e78bec457b35200f231d9f0b754c1 Mon Sep 17 00:00:00 2001 From: Angel Misevski Date: Thu, 16 Jan 2025 14:48:34 -0500 Subject: [PATCH 3/3] Fixup path handling in Kitfile generation on Windows To ensure generated Kitfiles are consistent, generate Kitfiles that always use the forward slash ('/') as a path separator. This means replacing slashes in windows-style paths with forward-slashes. --- pkg/lib/kitfile/generate.go | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/pkg/lib/kitfile/generate.go b/pkg/lib/kitfile/generate.go index 29e4662a9..53b363360 100644 --- a/pkg/lib/kitfile/generate.go +++ b/pkg/lib/kitfile/generate.go @@ -226,16 +226,18 @@ func addDirToKitfile(kitfile *artifact.KitFile, baseDir, dirPath string, d fs.Di directoryContents := [int(fileTypeUnknown) + 1][]string{} for _, entry := range entries { relPath := filepath.Join(dirPath, entry.Name()) + // We want to use unix-style paths (separated by '/') even if we're generating on Windows. + kitfileRelPath := filepath.ToSlash(relPath) if entry.IsDir() { // TODO: we can potentially recurse further here if we find we need to - directoryContents[int(fileTypeUnknown)] = append(directoryContents[int(fileTypeUnknown)], relPath) + directoryContents[int(fileTypeUnknown)] = append(directoryContents[int(fileTypeUnknown)], kitfileRelPath) continue } fileType := determineFileType(entry.Name()) if fileType == fileTypeModel { - modelFiles = append(modelFiles, relPath) + modelFiles = append(modelFiles, kitfileRelPath) } - directoryContents[int(fileType)] = append(directoryContents[int(fileType)], relPath) + directoryContents[int(fileType)] = append(directoryContents[int(fileType)], kitfileRelPath) } // Try to detect directories that contain e.g. only datasets so we can add them