test: wait for updater checker shutdown

app/ui/ui_test.go

@@ -14,6 +14,7 @@ import (
 	"strings"
 	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/ollama/ollama/api"
 	"github.com/ollama/ollama/app/store"
@@ -804,8 +805,13 @@ func TestSettingsToggleAutoUpdateOn_NoPendingUpdate_TriggersCheck(t *testing.T)
 	// Initialize the checkNow channel by starting (and immediately stopping) the checker
 	// so TriggerImmediateCheck doesn't panic on nil channel
 	ctx, cancel := context.WithCancel(t.Context())
-	upd.StartBackgroundUpdaterChecker(ctx, func(string) error { return nil })
-	defer cancel()
+	stopped := upd.StartBackgroundUpdaterChecker(ctx, func(string) error { return nil })
+	cancel()
+	select {
+	case <-stopped:
+	case <-time.After(2 * time.Second):
+		t.Fatal("background updater did not stop")
+	}
 
 	var notificationCalled atomic.Bool
 	server := &Server{

app/updater/updater.go

@@ -287,12 +287,22 @@ func (u *Updater) TriggerImmediateCheck() {
 	}
 }
 
-func (u *Updater) StartBackgroundUpdaterChecker(ctx context.Context, cb func(string) error) {
+func (u *Updater) StartBackgroundUpdaterChecker(ctx context.Context, cb func(string) error) <-chan struct{} {
+	stopped := make(chan struct{})
 	u.checkNow = make(chan struct{}, 1)
 	u.checkNow <- struct{}{} // Trigger first check after initial delay
 	go func() {
+		defer close(stopped)
+
 		// Don't blast an update message immediately after startup
-		time.Sleep(UpdateCheckInitialDelay)
+		initialDelay := time.NewTimer(UpdateCheckInitialDelay)
+		select {
+		case <-ctx.Done():
+			initialDelay.Stop()
+			slog.Debug("stopping background update checker")
+			return
+		case <-initialDelay.C:
+		}
 		slog.Info("beginning update checker", "interval", UpdateCheckInterval)
 		ticker := time.NewTicker(UpdateCheckInterval)
 		defer ticker.Stop()
@@ -341,4 +351,5 @@ func (u *Updater) StartBackgroundUpdaterChecker(ctx context.Context, cb func(str
 			}
 		}
 	}()
+	return stopped
 }

app/updater/updater_test.go

@@ -49,23 +49,26 @@ func TestIsNewReleaseAvailable(t *testing.T) {
 
 func TestBackgoundChecker(t *testing.T) {
 	UpdateStageDir = t.TempDir()
-	haveUpdate := false
-	verified := false
-	done := make(chan int)
+	haveUpdate := atomic.Bool{}
+	verified := atomic.Bool{}
+	updateDone := make(chan struct{}, 1)
 	cb := func(ver string) error {
-		haveUpdate = true
-		done <- 0
+		haveUpdate.Store(true)
+		select {
+		case updateDone <- struct{}{}:
+		default:
+		}
 		return nil
 	}
 	stallTimer := time.NewTimer(5 * time.Second)
+	defer stallTimer.Stop()
 	ctx, cancel := context.WithCancel(t.Context())
-	defer cancel()
 	UpdateCheckInitialDelay = 5 * time.Millisecond
 	UpdateCheckInterval = 5 * time.Millisecond
-	VerifyDownload = func() error {
-		verified = true
+	setVerifyDownload(t, func() error {
+		verified.Store(true)
 		return nil
-	}
+	})
 
 	var server *httptest.Server
 	server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -98,15 +101,16 @@ func TestBackgoundChecker(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	updater.StartBackgroundUpdaterChecker(ctx, cb)
+	stopped := updater.StartBackgroundUpdaterChecker(ctx, cb)
+	defer waitForBackgroundUpdater(t, cancel, stopped)
 	select {
 	case <-stallTimer.C:
 		t.Fatal("stalled")
-	case <-done:
-		if !haveUpdate {
+	case <-updateDone:
+		if !haveUpdate.Load() {
 			t.Fatal("no update received")
 		}
-		if !verified {
+		if !verified.Load() {
 			t.Fatal("unverified")
 		}
 	}
@@ -115,15 +119,14 @@ func TestBackgoundChecker(t *testing.T) {
 func TestAutoUpdateDisabledSkipsDownload(t *testing.T) {
 	UpdateStageDir = t.TempDir()
 	var downloadAttempted atomic.Bool
-	done := make(chan struct{})
+	var callbackCalled atomic.Bool
 
 	ctx, cancel := context.WithCancel(t.Context())
-	defer cancel()
 	UpdateCheckInitialDelay = 5 * time.Millisecond
 	UpdateCheckInterval = 5 * time.Millisecond
-	VerifyDownload = func() error {
+	setVerifyDownload(t, func() error {
 		return nil
-	}
+	})
 
 	var server *httptest.Server
 	server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -156,19 +159,22 @@ func TestAutoUpdateDisabledSkipsDownload(t *testing.T) {
 	}
 
 	cb := func(ver string) error {
-		t.Fatal("callback should not be called when auto-update is disabled")
+		callbackCalled.Store(true)
 		return nil
 	}
 
-	updater.StartBackgroundUpdaterChecker(ctx, cb)
+	stopped := updater.StartBackgroundUpdaterChecker(ctx, cb)
 
 	// Wait enough time for multiple check cycles
 	time.Sleep(50 * time.Millisecond)
-	close(done)
+	waitForBackgroundUpdater(t, cancel, stopped)
 
 	if downloadAttempted.Load() {
 		t.Fatal("download should not be attempted when auto-update is disabled")
 	}
+	if callbackCalled.Load() {
+		t.Fatal("callback should not be called when auto-update is disabled")
+	}
 }
 
 func TestAutoUpdateReenabledDownloadsUpdate(t *testing.T) {
@@ -177,12 +183,11 @@ func TestAutoUpdateReenabledDownloadsUpdate(t *testing.T) {
 	callbackCalled := make(chan struct{}, 1)
 
 	ctx, cancel := context.WithCancel(t.Context())
-	defer cancel()
 	UpdateCheckInitialDelay = 5 * time.Millisecond
 	UpdateCheckInterval = 5 * time.Millisecond
-	VerifyDownload = func() error {
+	setVerifyDownload(t, func() error {
 		return nil
-	}
+	})
 
 	var server *httptest.Server
 	server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -222,7 +227,8 @@ func TestAutoUpdateReenabledDownloadsUpdate(t *testing.T) {
 		return nil
 	}
 
-	upd.StartBackgroundUpdaterChecker(ctx, cb)
+	stopped := upd.StartBackgroundUpdaterChecker(ctx, cb)
+	defer waitForBackgroundUpdater(t, cancel, stopped)
 
 	// Wait for a few cycles with auto-update disabled - no download should happen
 	time.Sleep(50 * time.Millisecond)
@@ -254,9 +260,9 @@ func TestCancelOngoingDownload(t *testing.T) {
 	downloadCancelled := make(chan struct{})
 
 	ctx := t.Context()
-	VerifyDownload = func() error {
+	setVerifyDownload(t, func() error {
 		return nil
-	}
+	})
 
 	var server *httptest.Server
 	server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -320,13 +326,12 @@ func TestTriggerImmediateCheck(t *testing.T) {
 	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 {
+	setVerifyDownload(t, func() error {
 		return nil
-	}
+	})
 
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if r.URL.Path == "/update.json" {
@@ -349,7 +354,8 @@ func TestTriggerImmediateCheck(t *testing.T) {
 		return nil
 	}
 
-	updater.StartBackgroundUpdaterChecker(ctx, cb)
+	stopped := updater.StartBackgroundUpdaterChecker(ctx, cb)
+	defer waitForBackgroundUpdater(t, cancel, stopped)
 
 	// Wait for the initial check that fires after the initial delay
 	select {
@@ -375,3 +381,22 @@ func TestTriggerImmediateCheck(t *testing.T) {
 		t.Fatalf("TriggerImmediateCheck did not cause additional check: initial=%d, final=%d", initialCount, finalCount)
 	}
 }
+
+func waitForBackgroundUpdater(t *testing.T, cancel context.CancelFunc, stopped <-chan struct{}) {
+	t.Helper()
+	cancel()
+	select {
+	case <-stopped:
+	case <-time.After(2 * time.Second):
+		t.Fatal("background updater did not stop")
+	}
+}
+
+func setVerifyDownload(t *testing.T, fn func() error) {
+	t.Helper()
+	old := VerifyDownload
+	VerifyDownload = fn
+	t.Cleanup(func() {
+		VerifyDownload = old
+	})
+}