ollama/server/cache_test.go

package server

import (
	"os"
	"path/filepath"
	"testing"
	"time"
)

func TestModelCacheGet(t *testing.T) {
	testModel := &Model{
		Name:      "test-model:latest",
		ShortName: "test-model",
	}

	tests := []struct {
		name           string
		modelName      string
		setupFunc      func(t *testing.T, modelsDir string, cache *ModelCache) string // returns manifest path
		expectedModel  *Model
		expectedExists bool
	}{
		{
			name:      "cache hit - valid cached model",
			modelName: "test-model:latest",
			setupFunc: func(t *testing.T, modelsDir string, cache *ModelCache) string {
				createTestModel(t, modelsDir, "test-model", []Layer{
					{MediaType: "application/vnd.ollama.image.model", Digest: "sha256-abc123", Size: 1000},
				})

				manifestPath := filepath.Join(modelsDir, "manifests", "registry.ollama.ai", "library", "test-model", "latest")
				info, err := os.Stat(manifestPath)
				if err != nil {
					t.Fatal(err)
				}

				cache.cache["test-model:latest"] = &CachedModel{
					model:    testModel,
					modTime:  info.ModTime(),
					fileSize: info.Size(),
				}
				return manifestPath
			},
			expectedModel:  testModel,
			expectedExists: true,
		},
		{
			name:      "cache miss - no cached entry",
			modelName: "missing-model:latest",
			setupFunc: func(t *testing.T, modelsDir string, cache *ModelCache) string {
				createTestModel(t, modelsDir, "missing-model", []Layer{
					{MediaType: "application/vnd.ollama.image.model", Digest: "sha256-def456", Size: 2000},
				})
				return filepath.Join(modelsDir, "manifests", "registry.ollama.ai", "library", "missing-model", "latest")
			},
			expectedModel:  nil,
			expectedExists: false,
		},
		{
			name:      "cache stale - modification time changed",
			modelName: "stale-model:latest",
			setupFunc: func(t *testing.T, modelsDir string, cache *ModelCache) string {
				createTestModel(t, modelsDir, "stale-model", []Layer{
					{MediaType: "application/vnd.ollama.image.model", Digest: "sha256-ghi789", Size: 3000},
				})

				manifestPath := filepath.Join(modelsDir, "manifests", "registry.ollama.ai", "library", "stale-model", "latest")
				info, err := os.Stat(manifestPath)
				if err != nil {
					t.Fatal(err)
				}

				cache.cache["stale-model:latest"] = &CachedModel{
					model:    testModel,
					modTime:  info.ModTime().Add(-time.Hour), // Stale time
					fileSize: info.Size(),
				}
				return manifestPath
			},
			expectedModel:  nil,
			expectedExists: false,
		},
		{
			name:      "cache stale - file size changed",
			modelName: "stale-size-model:latest",
			setupFunc: func(t *testing.T, modelsDir string, cache *ModelCache) string {
				createTestModel(t, modelsDir, "stale-size-model", []Layer{
					{MediaType: "application/vnd.ollama.image.model", Digest: "sha256-jkl012", Size: 4000},
				})

				manifestPath := filepath.Join(modelsDir, "manifests", "registry.ollama.ai", "library", "stale-size-model", "latest")
				info, err := os.Stat(manifestPath)
				if err != nil {
					t.Fatal(err)
				}

				cache.cache["stale-size-model:latest"] = &CachedModel{
					model:    testModel,
					modTime:  info.ModTime(),
					fileSize: info.Size() + 100, // Different size
				}
				return manifestPath
			},
			expectedModel:  nil,
			expectedExists: false,
		},
		{
			name:      "manifest file does not exist",
			modelName: "nonexistent-model:latest",
			setupFunc: func(t *testing.T, modelsDir string, cache *ModelCache) string {
				// Add to cache but don't create manifest file
				cache.cache["nonexistent-model:latest"] = &CachedModel{
					model:    testModel,
					modTime:  time.Now(),
					fileSize: 100,
				}
				return "" // No manifest created
			},
			expectedModel:  nil,
			expectedExists: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			modelsDir := t.TempDir()
			t.Setenv("OLLAMA_MODELS", modelsDir)

			// Create fresh cache instance for each test
			cache := &ModelCache{
				cache: make(map[string]*CachedModel),
			}

			if tt.setupFunc != nil {
				tt.setupFunc(t, modelsDir, cache)
			}

			model, exists := cache.get(tt.modelName)

			if exists != tt.expectedExists {
				t.Errorf("get() exists = %v, want %v", exists, tt.expectedExists)
			}

			if tt.expectedModel != nil && model != tt.expectedModel {
				t.Errorf("get() model = %v, want %v", model, tt.expectedModel)
			}

			if tt.expectedModel == nil && model != nil {
				t.Errorf("get() model = %v, want nil", model)
			}
		})
	}
}

func TestModelCacheSet(t *testing.T) {
	testModel := &Model{
		Name:      "test-model:latest",
		ShortName: "test-model",
	}

	tests := []struct {
		name         string
		modelName    string
		model        *Model
		setupFunc    func(t *testing.T, modelsDir string) string // returns manifest path
		expectCached bool
	}{
		{
			name:      "successful cache set",
			modelName: "test-model:latest",
			model:     testModel,
			setupFunc: func(t *testing.T, modelsDir string) string {
				createTestModel(t, modelsDir, "test-model", []Layer{
					{MediaType: "application/vnd.ollama.image.model", Digest: "sha256-abc123", Size: 1000},
				})
				return filepath.Join(modelsDir, "manifests", "registry.ollama.ai", "library", "test-model", "latest")
			},
			expectCached: true,
		},
		{
			name:      "manifest file does not exist",
			modelName: "nonexistent-model:latest",
			model:     testModel,
			setupFunc: func(t *testing.T, modelsDir string) string {
				// Don't create manifest file
				return ""
			},
			expectCached: false,
		},
		{
			name:      "overwrite existing cache entry",
			modelName: "existing-model:latest",
			model:     testModel,
			setupFunc: func(t *testing.T, modelsDir string) string {
				createTestModel(t, modelsDir, "existing-model", []Layer{
					{MediaType: "application/vnd.ollama.image.model", Digest: "sha256-def456", Size: 2000},
				})
				return filepath.Join(modelsDir, "manifests", "registry.ollama.ai", "library", "existing-model", "latest")
			},
			expectCached: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			modelsDir := t.TempDir()
			t.Setenv("OLLAMA_MODELS", modelsDir)

			// Create fresh cache instance for each test
			cache := &ModelCache{
				cache: make(map[string]*CachedModel),
			}

			if tt.name == "overwrite existing cache entry" {
				cache.cache[tt.modelName] = &CachedModel{
					model:    &Model{Name: "old-model"},
					modTime:  time.Now().Add(-time.Hour),
					fileSize: 50,
				}
			}

			if tt.setupFunc != nil {
				tt.setupFunc(t, modelsDir)
			}

			cache.set(tt.modelName, tt.model)

			cached, exists := cache.cache[tt.modelName]

			if tt.expectCached {
				if !exists {
					t.Errorf("set() expected model to be cached, but it wasn't")
					return
				}

				if cached.model != tt.model {
					t.Errorf("set() cached model = %v, want %v", cached.model, tt.model)
				}

				// Verify file info is captured correctly if manifest exists
				expectedManifestPath := filepath.Join(modelsDir, "manifests", "registry.ollama.ai", "library", tt.modelName[:len(tt.modelName)-7], "latest")
				if info, err := os.Stat(expectedManifestPath); err == nil {
					if !cached.modTime.Equal(info.ModTime()) {
						t.Errorf("set() cached modTime = %v, want %v", cached.modTime, info.ModTime())
					}
					if cached.fileSize != info.Size() {
						t.Errorf("set() cached fileSize = %v, want %v", cached.fileSize, info.Size())
					}
				}
			} else {
				if exists {
					t.Errorf("set() expected model not to be cached, but it was")
				}
			}
		})
	}
}