Fix tests

Signed-off-by: David Gageot <david.gageot@docker.com>

Author: dgageot

pkg/agentfile/resolver_test.go

@@ -125,7 +125,13 @@ agents:
 
 	content1, err := os.ReadFile(resolved1)
 	require.NoError(t, err)
-	assert.Equal(t, agentContent, string(content1))
+	assert.Equal(t, `version: "1"
+agents:
+  root:
+    model: openai/gpt-4o
+    description: Test OCI agent
+    instruction: You are a test OCI agent
+`, string(content1))
 
 	// Expected filename based on OCI ref
 	expectedFilename := OciRefToFilename(ociRef)
@@ -454,3 +460,40 @@ func TestResolveAgentFile_ReplaceEmptyAliasWithActualFile(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, aliasedAgentFile, resolved)
 }
+
+func TestResolveAgentFile_OCIRef_HasAVersion(t *testing.T) {
+	storeDir := t.TempDir()
+	t.Setenv("CAGENT_CONTENT_STORE", storeDir)
+
+	store, err := content.NewStore()
+	require.NoError(t, err)
+
+	agentContent := `agents:
+  root:
+    model: auto
+    description: Test OCI agent
+    instruction: You are a test OCI agent
+`
+	agentFile := filepath.Join(t.TempDir(), "oci-agent.yaml")
+	require.NoError(t, os.WriteFile(agentFile, []byte(agentContent), 0o644))
+
+	// Package as OCI artifact
+	ociRef := "test.registry.io/myorg/testagent:v1"
+	_, err = oci.PackageFileAsOCIToStore(t.Context(), agentFile, ociRef, store)
+	require.NoError(t, err)
+
+	// First resolution
+	resolved, err := Resolve(t.Context(), nil, ociRef)
+	require.NoError(t, err)
+	assert.NotEmpty(t, resolved)
+
+	content, err := os.ReadFile(resolved)
+	require.NoError(t, err)
+	assert.Equal(t, `version: "2"
+agents:
+  root:
+    model: auto
+    description: Test OCI agent
+    instruction: You are a test OCI agent
+`, string(content))
+}

pkg/oci/package.go

@@ -44,11 +44,19 @@ func PackageFileAsOCIToStore(ctx context.Context, filePath, artifactRef string,
 		return "", fmt.Errorf("loading config: %w", err)
 	}
 
-	// Make sure we push a yaml with the latest version
-	cfg.Version = latest.Version
-	data, err = yaml.MarshalWithOptions(cfg, yaml.Indent(2))
-	if err != nil {
-		return "", fmt.Errorf("marshaling config: %w", err)
+	var raw struct {
+		Version string `yaml:"version,omitempty"`
+	}
+	if err := yaml.UnmarshalWithOptions(data, &raw); err != nil {
+		return "", fmt.Errorf("looking for version in config file\n%s", yaml.FormatError(err, true, true))
+	}
+	// Make sure we push a yaml with a version (Use latest if missing)
+	if raw.Version == "" {
+		cfg.Version = latest.Version
+		data, err = yaml.MarshalWithOptions(cfg, yaml.Indent(2))
+		if err != nil {
+			return "", fmt.Errorf("marshaling config: %w", err)
+		}
 	}
 
 	// Prepare OCI annotations

pkg/oci/package_test.go

@@ -13,9 +13,11 @@ import (
 
 func TestPackageFileAsOCIToStore(t *testing.T) {
 	testFile := filepath.Join(t.TempDir(), "test.yaml")
-	testContent := `name: test-app
-version: v1.0.0
-description: "Test application"
+	testContent := `version: "2"
+agents:
+  root:
+    model: auto
+    description: A helpful AI assistant
 `
 	require.NoError(t, os.WriteFile(testFile, []byte(testContent), 0o644))
 	store, err := content.NewStore(content.WithBaseDir(t.TempDir()))
@@ -24,9 +26,7 @@ description: "Test application"
 	tag := "test-app:v1.0.0"
 	digest, err := PackageFileAsOCIToStore(t.Context(), testFile, tag, store)
 	require.NoError(t, err)
-
 	assert.NotEmpty(t, digest)
-
 	t.Cleanup(func() {
 		if err := store.DeleteArtifact(digest); err != nil {
 			t.Logf("Failed to clean up artifact: %v", err)
@@ -35,7 +35,6 @@ description: "Test application"
 
 	img, err := store.GetArtifactImage(tag)
 	require.NoError(t, err)
-
 	assert.NotNil(t, img)
 
 	metadata, err := store.GetArtifactMetadata(tag)
@@ -67,61 +66,3 @@ func TestPackageFileAsOCIToStoreInvalidTag(t *testing.T) {
 	_, err = PackageFileAsOCIToStore(t.Context(), testFile, "", store)
 	require.Error(t, err)
 }
-
-func TestPackageFileAsOCIToStoreDifferentFileTypes(t *testing.T) {
-	testCases := []struct {
-		name     string
-		filename string
-		content  string
-		tag      string
-	}{
-		{
-			name:     "yaml file",
-			filename: "config.yaml",
-			content:  "key: value\nother: data",
-			tag:      "config:yaml",
-		},
-		{
-			name:     "json file",
-			filename: "data.json",
-			content:  `{"key": "value", "number": 42}`,
-			tag:      "data:json",
-		},
-		{
-			name:     "text file",
-			filename: "readme.txt",
-			content:  "This is a simple text file\nwith multiple lines",
-			tag:      "readme:txt",
-		},
-	}
-
-	store, err := content.NewStore(content.WithBaseDir(t.TempDir()))
-	require.NoError(t, err)
-
-	var digests []string
-
-	for _, tc := range testCases {
-		t.Run(tc.name, func(t *testing.T) {
-			testFile := filepath.Join(t.TempDir(), tc.filename)
-			require.NoError(t, os.WriteFile(testFile, []byte(tc.content), 0o644))
-
-			// Package the file as OCI artifact
-			digest, err := PackageFileAsOCIToStore(t.Context(), testFile, tc.tag, store)
-			require.NoError(t, err)
-
-			digests = append(digests, digest)
-
-			img, err := store.GetArtifactImage(tc.tag)
-			require.NoError(t, err)
-			assert.NotNil(t, img)
-		})
-	}
-
-	t.Cleanup(func() {
-		for _, digest := range digests {
-			if err := store.DeleteArtifact(digest); err != nil {
-				t.Logf("Failed to clean up artifact %s: %v", digest, err)
-			}
-		}
-	})
-}