ollama/app/updater/updater_test.go

//go:build windows || darwin

package updater

import (
	"archive/zip"
	"bytes"
	"context"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"sync/atomic"
	"testing"
	"time"

	"github.com/ollama/ollama/app/store"
)

func TestIsNewReleaseAvailable(t *testing.T) {
	slog.SetLogLoggerLevel(slog.LevelDebug)
	var server *httptest.Server
	server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/update.json" {
			w.Write([]byte(
				fmt.Sprintf(`{"version": "9.9.9", "url": "%s"}`,
					server.URL+"/9.9.9/"+Installer)))
			// TODO - wire up the redirects to mimic real behavior
		} else {
			slog.Debug("unexpected request", "url", r.URL)
		}
	}))
	defer server.Close()
	slog.Debug("server", "url", server.URL)

	updater := &Updater{Store: &store.Store{DBPath: filepath.Join(t.TempDir(), "test.db")}}
	defer updater.Store.Close() // Ensure database is closed
	UpdateCheckURLBase = server.URL + "/update.json"
	updatePresent, resp := updater.checkForUpdate(t.Context())
	if !updatePresent {
		t.Fatal("expected update to be available")
	}
	if resp.UpdateVersion != "9.9.9" {
		t.Fatal("unexpected response", "url", resp.UpdateURL, "version", resp.UpdateVersion)
	}
}

func TestBackgoundChecker(t *testing.T) {
	UpdateStageDir = t.TempDir()
	haveUpdate := false
	verified := false
	done := make(chan int)
	cb := func(ver string) error {
		haveUpdate = true
		done <- 0
		return nil
	}
	stallTimer := time.NewTimer(5 * time.Second)
	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()
	UpdateCheckInitialDelay = 5 * time.Millisecond
	UpdateCheckInterval = 5 * time.Millisecond
	VerifyDownload = func() error {
		verified = true
		return nil
	}

	var server *httptest.Server
	server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/update.json" {
			w.Write([]byte(
				fmt.Sprintf(`{"version": "9.9.9", "url": "%s"}`,
					server.URL+"/9.9.9/"+Installer)))
			// TODO - wire up the redirects to mimic real behavior
		} else if r.URL.Path == "/9.9.9/"+Installer {
			buf := &bytes.Buffer{}
			zw := zip.NewWriter(buf)
			zw.Close()
			io.Copy(w, buf)
		} else {
			slog.Debug("unexpected request", "url", r.URL)
		}
	}))
	defer server.Close()
	UpdateCheckURLBase = server.URL + "/update.json"

	updater := &Updater{Store: &store.Store{DBPath: filepath.Join(t.TempDir(), "test.db")}}
	defer updater.Store.Close()

	settings, err := updater.Store.Settings()
	if err != nil {
		t.Fatal(err)
	}
	settings.AutoUpdateEnabled = true
	if err := updater.Store.SetSettings(settings); err != nil {
		t.Fatal(err)
	}

	updater.StartBackgroundUpdaterChecker(ctx, cb)
	select {
	case <-stallTimer.C:
		t.Fatal("stalled")
	case <-done:
		if !haveUpdate {
			t.Fatal("no update received")
		}
		if !verified {
			t.Fatal("unverified")
		}
	}
}

func TestAutoUpdateDisabledSkipsDownload(t *testing.T) {
	UpdateStageDir = t.TempDir()
	var downloadAttempted atomic.Bool
	done := make(chan struct{})

	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()
	UpdateCheckInitialDelay = 5 * time.Millisecond
	UpdateCheckInterval = 5 * time.Millisecond
	VerifyDownload = func() error {
		return nil
	}

	var server *httptest.Server
	server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/update.json" {
			w.Write([]byte(
				fmt.Sprintf(`{"version": "9.9.9", "url": "%s"}`,
					server.URL+"/9.9.9/"+Installer)))
		} else if r.URL.Path == "/9.9.9/"+Installer {
			downloadAttempted.Store(true)
			buf := &bytes.Buffer{}
			zw := zip.NewWriter(buf)
			zw.Close()
			io.Copy(w, buf)
		}
	}))
	defer server.Close()
	UpdateCheckURLBase = server.URL + "/update.json"

	updater := &Updater{Store: &store.Store{DBPath: filepath.Join(t.TempDir(), "test.db")}}
	defer updater.Store.Close()

	// Ensure auto-update is disabled
	settings, err := updater.Store.Settings()
	if err != nil {
		t.Fatal(err)
	}
	settings.AutoUpdateEnabled = false
	if err := updater.Store.SetSettings(settings); err != nil {
		t.Fatal(err)
	}

	cb := func(ver string) error {
		t.Fatal("callback should not be called when auto-update is disabled")
		return nil
	}

	updater.StartBackgroundUpdaterChecker(ctx, cb)

	// Wait enough time for multiple check cycles
	time.Sleep(50 * time.Millisecond)
	close(done)

	if downloadAttempted.Load() {
		t.Fatal("download should not be attempted when auto-update is disabled")
	}
}

func TestAutoUpdateReenabledDownloadsUpdate(t *testing.T) {
	UpdateStageDir = t.TempDir()
	var downloadAttempted atomic.Bool
	callbackCalled := make(chan struct{}, 1)

	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()
	UpdateCheckInitialDelay = 5 * time.Millisecond
	UpdateCheckInterval = 5 * time.Millisecond
	VerifyDownload = func() error {
		return nil
	}

	var server *httptest.Server
	server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/update.json" {
			w.Write([]byte(
				fmt.Sprintf(`{"version": "9.9.9", "url": "%s"}`,
					server.URL+"/9.9.9/"+Installer)))
		} else if r.URL.Path == "/9.9.9/"+Installer {
			downloadAttempted.Store(true)
			buf := &bytes.Buffer{}
			zw := zip.NewWriter(buf)
			zw.Close()
			io.Copy(w, buf)
		}
	}))
	defer server.Close()
	UpdateCheckURLBase = server.URL + "/update.json"

	upd := &Updater{Store: &store.Store{DBPath: filepath.Join(t.TempDir(), "test.db")}}
	defer upd.Store.Close()

	// Start with auto-update disabled
	settings, err := upd.Store.Settings()
	if err != nil {
		t.Fatal(err)
	}
	settings.AutoUpdateEnabled = false
	if err := upd.Store.SetSettings(settings); err != nil {
		t.Fatal(err)
	}

	cb := func(ver string) error {
		select {
		case callbackCalled <- struct{}{}:
		default:
		}
		return nil
	}

	upd.StartBackgroundUpdaterChecker(ctx, cb)

	// Wait for a few cycles with auto-update disabled - no download should happen
	time.Sleep(50 * time.Millisecond)
	if downloadAttempted.Load() {
		t.Fatal("download should not happen while auto-update is disabled")
	}

	// Re-enable auto-update
	settings.AutoUpdateEnabled = true
	if err := upd.Store.SetSettings(settings); err != nil {
		t.Fatal(err)
	}

	// Wait for the checker to pick it up and download
	select {
	case <-callbackCalled:
		// Success: download happened and callback was called after re-enabling
		if !downloadAttempted.Load() {
			t.Fatal("expected download to be attempted after re-enabling")
		}
	case <-time.After(5 * time.Second):
		t.Fatal("expected download and callback after re-enabling auto-update")
	}
}

func TestCancelOngoingDownload(t *testing.T) {
	UpdateStageDir = t.TempDir()
	downloadStarted := make(chan struct{})
	downloadCancelled := make(chan struct{})

	ctx := t.Context()
	VerifyDownload = func() error {
		return nil
	}

	var server *httptest.Server
	server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/update.json" {
			w.Write([]byte(
				fmt.Sprintf(`{"version": "9.9.9", "url": "%s"}`,
					server.URL+"/9.9.9/"+Installer)))
		} else if r.URL.Path == "/9.9.9/"+Installer {
			if r.Method == http.MethodHead {
				w.Header().Set("Content-Length", "1000000")
				w.WriteHeader(http.StatusOK)
				return
			}
			// Signal that download has started
			close(downloadStarted)
			// Wait for cancellation or timeout
			select {
			case <-r.Context().Done():
				close(downloadCancelled)
				return
			case <-time.After(5 * time.Second):
				t.Error("download was not cancelled in time")
			}
		}
	}))
	defer server.Close()
	UpdateCheckURLBase = server.URL + "/update.json"

	updater := &Updater{Store: &store.Store{DBPath: filepath.Join(t.TempDir(), "test.db")}}
	defer updater.Store.Close()

	_, resp := updater.checkForUpdate(ctx)

	// Start download in goroutine
	go func() {
		_ = updater.DownloadNewRelease(ctx, resp)
	}()

	// Wait for download to start
	select {
	case <-downloadStarted:
	case <-time.After(2 * time.Second):
		t.Fatal("download did not start in time")
	}

	// Cancel the download
	updater.CancelOngoingDownload()

	// Verify cancellation was received
	select {
	case <-downloadCancelled:
		// Success
	case <-time.After(2 * time.Second):
		t.Fatal("download cancellation was not received by server")
	}
}

func TestTriggerImmediateCheck(t *testing.T) {
	UpdateStageDir = t.TempDir()
	checkCount := atomic.Int32{}
	checkDone := make(chan struct{}, 10)

	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()
	// Set a very long interval so only TriggerImmediateCheck causes checks
	UpdateCheckInitialDelay = 1 * time.Millisecond
	UpdateCheckInterval = 1 * time.Hour
	VerifyDownload = func() error {
		return nil
	}

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/update.json" {
			checkCount.Add(1)
			select {
			case checkDone <- struct{}{}:
			default:
			}
			// Return no update available
			w.WriteHeader(http.StatusNoContent)
		}
	}))
	defer server.Close()
	UpdateCheckURLBase = server.URL + "/update.json"

	updater := &Updater{Store: &store.Store{DBPath: filepath.Join(t.TempDir(), "test.db")}}
	defer updater.Store.Close()

	cb := func(ver string) error {
		return nil
	}

	updater.StartBackgroundUpdaterChecker(ctx, cb)

	// Wait for the initial check that fires after the initial delay
	select {
	case <-checkDone:
	case <-time.After(2 * time.Second):
		t.Fatal("initial check did not happen")
	}

	initialCount := checkCount.Load()

	// Trigger immediate check
	updater.TriggerImmediateCheck()

	// Wait for the triggered check
	select {
	case <-checkDone:
	case <-time.After(2 * time.Second):
		t.Fatal("triggered check did not happen")
	}

	finalCount := checkCount.Load()
	if finalCount <= initialCount {
		t.Fatalf("TriggerImmediateCheck did not cause additional check: initial=%d, final=%d", initialCount, finalCount)
	}
}