Add tests for generating Kitfiles.

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

testing/generate-kitfile_test.go

@@ -0,0 +1,87 @@
+// 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) {
+	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 := setupTempDir(t)
+
+			// 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
+			}
+			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
+			}
+			assert.Equal(t, tt.ExpectedKitfile, actualKitfile, "Generated Kitfile should match")
+		})
+	}
+}

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")

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)

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)

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

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"

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

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

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

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