Add test to ensure short IDs are properly resolved

Author: ilopezluna

pkg/distribution/distribution/normalize_test.go

@@ -1,8 +1,14 @@
 package distribution
 
 import (
+	"context"
+	"io"
+	"path/filepath"
+	"strings"
 	"testing"
 
+	"github.com/docker/model-runner/pkg/distribution/builder"
+	"github.com/docker/model-runner/pkg/distribution/tarball"
 	"github.com/sirupsen/logrus"
 )
 
@@ -276,6 +282,54 @@ func TestLooksLikeDigest(t *testing.T) {
 	}
 }
 
+func TestNormalizeModelNameWithIDResolution(t *testing.T) {
+	// Create a client with a temporary store
+	client, cleanup := createTestClient(t)
+	defer cleanup()
+
+	// Load a test model to get a real ID
+	testGGUFFile := filepath.Join("..", "assets", "dummy.gguf")
+	modelID := loadTestModel(t, client, testGGUFFile)
+
+	// Extract the short ID (12 hex chars after "sha256:")
+	if !strings.HasPrefix(modelID, "sha256:") {
+		t.Fatalf("Expected model ID to start with 'sha256:', got: %s", modelID)
+	}
+	shortID := modelID[7:19] // Extract 12 chars after "sha256:"
+	fullHex := strings.TrimPrefix(modelID, "sha256:")
+
+	tests := []struct {
+		name     string
+		input    string
+		expected string
+	}{
+		{
+			name:     "short ID resolves to full ID",
+			input:    shortID,
+			expected: modelID,
+		},
+		{
+			name:     "full hex (without sha256:) resolves to full ID",
+			input:    fullHex,
+			expected: modelID,
+		},
+		{
+			name:     "full digest (with sha256:) returns as-is",
+			input:    modelID,
+			expected: modelID,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := client.normalizeModelName(tt.input)
+			if result != tt.expected {
+				t.Errorf("normalizeModelName(%q) = %q, want %q", tt.input, result, tt.expected)
+			}
+		})
+	}
+}
+
 // Helper function to create a test client with temp store
 func createTestClient(t *testing.T) (*Client, func()) {
 	t.Helper()
@@ -300,3 +354,43 @@ func createTestClient(t *testing.T) (*Client, func()) {
 
 	return client, cleanup
 }
+
+// Helper function to load a test model and return its ID
+func loadTestModel(t *testing.T, client *Client, ggufPath string) string {
+	t.Helper()
+
+	// Load model using LoadModel
+	pr, pw := io.Pipe()
+	target, err := tarball.NewTarget(pw)
+	if err != nil {
+		t.Fatalf("Failed to create target: %v", err)
+	}
+
+	done := make(chan error)
+	var id string
+	go func() {
+		var err error
+		id, err = client.LoadModel(pr, nil)
+		done <- err
+	}()
+
+	bldr, err := builder.FromGGUF(ggufPath)
+	if err != nil {
+		t.Fatalf("Failed to create builder from GGUF: %v", err)
+	}
+
+	ctx := context.Background()
+	if err := bldr.Build(ctx, target, nil); err != nil {
+		t.Fatalf("Failed to build model: %v", err)
+	}
+
+	if err := <-done; err != nil {
+		t.Fatalf("Failed to load model: %v", err)
+	}
+
+	if id == "" {
+		t.Fatal("Model ID is empty")
+	}
+
+	return id
+}