Ctrl K
ComunidadGuías…

internal/hub/download_test.go

Rama main · réplica de lectura del repositorio original.

Archivo original ↗Preguntar a KiBot
package hub

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

// fakeSource is a Source backed by an httptest server, so the download manager
// can be exercised end to end without touching Hugging Face.
type fakeSource struct {
	base  string
	files []File
}

func (f *fakeSource) ID() string   { return "fake" }
func (f *fakeSource) Name() string { return "Fake" }
func (f *fakeSource) Search(context.Context, Query) (Page, error) {
	return Page{Models: []Model{{ID: "o/r", Source: "fake"}}}, nil
}
func (f *fakeSource) Detail(_ context.Context, repo string) (ModelDetail, error) {
	return ModelDetail{Model: Model{ID: repo, Source: "fake"}, Files: f.files}, nil
}
func (f *fakeSource) FileURL(repo, path string) (string, error) {
	return f.base + "/" + repo + "/" + path, nil
}
func (f *fakeSource) CheckURL(raw string) error {
	if !strings.HasPrefix(raw, f.base) {
		return fmt.Errorf("refusing %s", raw)
	}
	return nil
}
func (f *fakeSource) Authorize(*http.Request) {}

// serveBlob answers a GET with Range support, optionally cutting the response
// short after cutAfter bytes to simulate a dropped connection.
func serveBlob(w http.ResponseWriter, r *http.Request, blob []byte, cutAfter int) {
	start := 0
	if rng := r.Header.Get("Range"); strings.HasPrefix(rng, "bytes=") {
		n, err := strconv.Atoi(strings.TrimSuffix(strings.TrimPrefix(rng, "bytes="), "-"))
		if err == nil {
			start = n
		}
		if start >= len(blob) {
			w.WriteHeader(http.StatusRequestedRangeNotSatisfiable)
			return
		}
		w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, len(blob)-1, len(blob)))
		w.WriteHeader(http.StatusPartialContent)
	}
	body := blob[start:]
	if cutAfter > 0 && cutAfter < len(body) {
		body = body[:cutAfter]
		w.Write(body)
		// No Content-Length was set, so a short write plus a flush and a panic
		// closes the connection mid-body the way a dropped line does.
		if fl, ok := w.(http.Flusher); ok {
			fl.Flush()
		}
		panic(http.ErrAbortHandler)
	}
	w.Write(body)
}

// discardLogger silences httptest's "connection hijacked/aborted" noise for the
// tests that deliberately drop a connection.
func discardLogger() *log.Logger { return log.New(io.Discard, "", 0) }

func waitJob(t *testing.T, m *Manager, id string) Job {
	t.Helper()
	deadline := time.Now().Add(20 * time.Second)
	for time.Now().Before(deadline) {
		j, ok := m.Job(id)
		if ok && j.Done() {
			return j
		}
		time.Sleep(10 * time.Millisecond)
	}
	j, _ := m.Job(id)
	t.Fatalf("job %s did not finish, stuck in phase %q", id, j.Phase)
	return Job{}
}

func blobOf(n int) []byte {
	b := make([]byte, n)
	for i := range b {
		b[i] = byte(i % 251)
	}
	return b
}

func TestManager_DownloadsEveryFile(t *testing.T) {
	a, b := blobOf(4096), blobOf(1500)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasSuffix(r.URL.Path, "b.gguf") {
			serveBlob(w, r, b, 0)
			return
		}
		serveBlob(w, r, a, 0)
	}))
	defer srv.Close()

	root := t.TempDir()
	src := &fakeSource{base: srv.URL, files: []File{
		{Path: "a-00001-of-00002.gguf", SizeBytes: int64(len(a))},
		{Path: "sub/b.gguf", SizeBytes: int64(len(b))},
	}}
	var completed atomic.Int32
	m := NewManager(func() string { return root }, nil, src)
	m.OnComplete = func(Job) error { completed.Add(1); return nil }

	id, err := m.Start(context.Background(), StartRequest{
		Source: "fake", Repo: "o/r", Files: []string{"a-00001-of-00002.gguf", "sub/b.gguf"},
	})
	if err != nil {
		t.Fatal(err)
	}
	job := waitJob(t, m, id)
	if job.Phase != PhaseDone {
		t.Fatalf("phase = %q, err = %q", job.Phase, job.Err)
	}
	dir := filepath.Join(root, "o", "r")
	if got, err := os.ReadFile(filepath.Join(dir, "a-00001-of-00002.gguf")); err != nil || !bytes.Equal(got, a) {
		t.Errorf("shard 1 mismatch: err %v, %d bytes", err, len(got))
	}
	if got, err := os.ReadFile(filepath.Join(dir, "sub", "b.gguf")); err != nil || !bytes.Equal(got, b) {
		t.Errorf("nested file mismatch: err %v, %d bytes", err, len(got))
	}
	if job.Downloaded != job.Total {
		t.Errorf("progress ended at %d/%d", job.Downloaded, job.Total)
	}
	if completed.Load() != 1 {
		t.Errorf("OnComplete ran %d times, want 1", completed.Load())
	}
	// Nothing partial may survive: model discovery walks this folder.
	if _, err := os.Stat(filepath.Join(dir, "a-00001-of-00002.gguf"+partSuffix)); !os.IsNotExist(err) {
		t.Error(".part file left behind after a successful download")
	}
	// Nor may the journal: there is nothing left to resume.
	if _, err := os.Stat(filepath.Join(dir, journalName)); !os.IsNotExist(err) {
		t.Error("the download journal survived a finished job")
	}
}

func TestManager_ResumesAfterDroppedConnection(t *testing.T) {
	blob := blobOf(8192)
	var hits atomic.Int32
	var sawRange atomic.Bool
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Range") != "" {
			sawRange.Store(true)
		}
		// Cut the first attempt short; serve the rest normally.
		cut := 0
		if hits.Add(1) == 1 {
			cut = 3000
		}
		serveBlob(w, r, blob, cut)
	}))
	defer srv.Close()
	srv.Config.ErrorLog = discardLogger()

	root := t.TempDir()
	src := &fakeSource{base: srv.URL, files: []File{{Path: "m.gguf", SizeBytes: int64(len(blob))}}}
	m := NewManager(func() string { return root }, nil, src)

	id, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"m.gguf"}})
	if err != nil {
		t.Fatal(err)
	}
	job := waitJob(t, m, id)
	if job.Phase != PhaseDone {
		t.Fatalf("phase = %q, err = %q", job.Phase, job.Err)
	}
	if !sawRange.Load() {
		t.Error("retry did not send a Range header — the transfer restarted from zero")
	}
	got, err := os.ReadFile(filepath.Join(root, "o", "r", "m.gguf"))
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(got, blob) {
		t.Errorf("resumed file is %d bytes and does not match the source", len(got))
	}
}

func TestManager_SkipsFileAlreadyOnDisk(t *testing.T) {
	blob := blobOf(2048)
	var hits atomic.Int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		serveBlob(w, r, blob, 0)
	}))
	defer srv.Close()

	root := t.TempDir()
	dir := filepath.Join(root, "o", "r")
	if err := os.MkdirAll(dir, 0o755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(dir, "m.gguf"), blob, 0o644); err != nil {
		t.Fatal(err)
	}
	src := &fakeSource{base: srv.URL, files: []File{{Path: "m.gguf", SizeBytes: int64(len(blob))}}}
	m := NewManager(func() string { return root }, nil, src)

	id, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"m.gguf"}})
	if err != nil {
		t.Fatal(err)
	}
	job := waitJob(t, m, id)
	if job.Phase != PhaseDone {
		t.Fatalf("phase = %q, err = %q", job.Phase, job.Err)
	}
	if hits.Load() != 0 {
		t.Errorf("refetched a file already on disk (%d requests)", hits.Load())
	}
	if !job.Files[0].Skipped || job.Downloaded != job.Total {
		t.Errorf("skipped file not accounted for: %+v (%d/%d)", job.Files[0], job.Downloaded, job.Total)
	}
}

func TestManager_GatedRepoFailsFastAsAuthError(t *testing.T) {
	var hits atomic.Int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		w.WriteHeader(http.StatusForbidden)
		w.Write([]byte("access to model o/r is restricted"))
	}))
	defer srv.Close()

	root := t.TempDir()
	src := &fakeSource{base: srv.URL, files: []File{{Path: "m.gguf", SizeBytes: 10}}}
	m := NewManager(func() string { return root }, nil, src)

	id, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"m.gguf"}})
	if err != nil {
		t.Fatal(err)
	}
	job := waitJob(t, m, id)
	if job.Phase != PhaseError || !job.Gated {
		t.Fatalf("phase = %q gated = %v, want error + gated", job.Phase, job.Gated)
	}
	// A license wall is terminal: retrying it five times only delays telling
	// the user to accept the license.
	if hits.Load() != 1 {
		t.Errorf("retried a 403 %d times", hits.Load())
	}
	if !strings.Contains(job.Err, "accept the license") {
		t.Errorf("error does not tell the user what to do: %q", job.Err)
	}
}

// heldServer serves the first `head` bytes of blob and then holds the
// connection open until the test lets go, so a transfer can be caught mid-flight
// and interrupted deterministically.
func heldServer(t *testing.T, blob []byte, head int) (*httptest.Server, func()) {
	t.Helper()
	release := make(chan struct{})
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write(blob[:head])
		if fl, ok := w.(http.Flusher); ok {
			fl.Flush()
		}
		<-release
	}))
	srv.Config.ErrorLog = discardLogger()
	return srv, func() { close(release); srv.Close() }
}

func waitPhase(t *testing.T, m *Manager, id, phase string) Job {
	t.Helper()
	deadline := time.Now().Add(20 * time.Second)
	for time.Now().Before(deadline) {
		if j, ok := m.Job(id); ok && j.Phase == phase {
			return j
		}
		time.Sleep(10 * time.Millisecond)
	}
	j, _ := m.Job(id)
	t.Fatalf("job %s never reached %q, stuck in %q", id, phase, j.Phase)
	return Job{}
}

// waitPartial waits until there are real bytes in the .part file. Neither of
// the job's own signals can say that: PhaseDownloading is set before the
// connection is even made (so the UI says "downloading" rather than sitting on
// "resolving" through DNS and TLS), and Downloaded is reported at most once per
// progressEvery, so a transfer smaller than a MiB never bumps it off zero. The
// file on disk is the thing this test is actually about.
func waitPartial(t *testing.T, part string) {
	t.Helper()
	deadline := time.Now().Add(20 * time.Second)
	for time.Now().Before(deadline) {
		if st, err := os.Stat(part); err == nil && st.Size() > 0 {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatalf("no bytes landed in %s", part)
}

func TestManager_PauseKeepsPartialAndResumeFinishes(t *testing.T) {
	blob := blobOf(1 << 20)
	srv, stop := heldServer(t, blob, 4096)

	root := t.TempDir()
	src := &fakeSource{base: srv.URL, files: []File{{Path: "m.gguf", SizeBytes: int64(len(blob))}}}
	m := NewManager(func() string { return root }, nil, src)

	id, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"m.gguf"}})
	if err != nil {
		t.Fatal(err)
	}
	part := filepath.Join(root, "o", "r", "m.gguf"+partSuffix)
	waitPartial(t, part)
	if err := m.Pause(id); err != nil {
		t.Fatal(err)
	}
	waitPhase(t, m, id, PhasePaused)

	// The partial must survive — that is the entire difference from cancel.
	st, err := os.Stat(part)
	if err != nil {
		t.Fatalf("pause discarded the partial file: %v", err)
	}
	if st.Size() == 0 {
		t.Fatal("pause left an empty partial")
	}
	if _, err := os.Stat(filepath.Join(root, "o", "r", "m.gguf")); !os.IsNotExist(err) {
		t.Error("a paused download left a file that looks complete")
	}

	// Resume against a server that now answers in full: the job keeps its id and
	// finishes from the bytes already on disk.
	stop()
	full := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		serveBlob(w, r, blob, 0)
	}))
	defer full.Close()
	src.base = full.URL

	if err := m.Resume(id); err != nil {
		t.Fatal(err)
	}
	job := waitJob(t, m, id)
	if job.Phase != PhaseDone {
		t.Fatalf("phase = %q, err = %q", job.Phase, job.Err)
	}
	got, err := os.ReadFile(filepath.Join(root, "o", "r", "m.gguf"))
	if err != nil || !bytes.Equal(got, blob) {
		t.Errorf("resumed file mismatch: err %v, %d bytes", err, len(got))
	}
}

func TestManager_CancelDiscardsPartial(t *testing.T) {
	blob := blobOf(1 << 20)
	srv, stop := heldServer(t, blob, 4096)
	defer stop()

	root := t.TempDir()
	src := &fakeSource{base: srv.URL, files: []File{{Path: "m.gguf", SizeBytes: int64(len(blob))}}}
	m := NewManager(func() string { return root }, nil, src)

	id, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"m.gguf"}})
	if err != nil {
		t.Fatal(err)
	}
	waitPhase(t, m, id, PhaseDownloading)
	if err := m.Cancel(id); err != nil {
		t.Fatal(err)
	}
	job := waitJob(t, m, id)
	if job.Phase != PhaseCanceled {
		t.Fatalf("phase = %q, want canceled (err %q)", job.Phase, job.Err)
	}
	// Cancel means discard: an abandoned .part is a multi-GB orphan nothing ever
	// mentions again.
	if _, err := os.Stat(filepath.Join(root, "o", "r", "m.gguf"+partSuffix)); !os.IsNotExist(err) {
		t.Errorf("cancel kept the partial file: %v", err)
	}
	if _, err := os.Stat(filepath.Join(root, "o", "r", "m.gguf")); !os.IsNotExist(err) {
		t.Error("a canceled download left a file that looks complete")
	}
	// The empty folders it created go with it.
	if _, err := os.Stat(filepath.Join(root, "o")); !os.IsNotExist(err) {
		t.Error("cancel left the publisher folder behind")
	}
}

func TestManager_CancelDiscardsFilesTheJobFinished(t *testing.T) {
	// Two files: one lands complete, the second is still running when the job is
	// canceled. A lone shard is not a model, so the finished one must go too —
	// but a file that was already there before the job (Skipped) must not.
	done, running := blobOf(2048), blobOf(1<<20)
	release := make(chan struct{})
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasSuffix(r.URL.Path, "a.gguf") {
			serveBlob(w, r, done, 0)
			return
		}
		w.Write(running[:4096])
		if fl, ok := w.(http.Flusher); ok {
			fl.Flush()
		}
		<-release
	}))
	defer srv.Close()
	defer close(release)
	srv.Config.ErrorLog = discardLogger()

	root := t.TempDir()
	dir := filepath.Join(root, "o", "r")
	if err := os.MkdirAll(dir, 0o755); err != nil {
		t.Fatal(err)
	}
	pre := blobOf(512)
	if err := os.WriteFile(filepath.Join(dir, "pre.gguf"), pre, 0o644); err != nil {
		t.Fatal(err)
	}
	src := &fakeSource{base: srv.URL, files: []File{
		{Path: "pre.gguf", SizeBytes: int64(len(pre))},
		{Path: "a.gguf", SizeBytes: int64(len(done))},
		{Path: "b.gguf", SizeBytes: int64(len(running))},
	}}
	m := NewManager(func() string { return root }, nil, src)

	id, err := m.Start(context.Background(), StartRequest{
		Source: "fake", Repo: "o/r", Files: []string{"pre.gguf", "a.gguf", "b.gguf"},
	})
	if err != nil {
		t.Fatal(err)
	}
	// b.gguf is the one that hangs, so wait until a.gguf has landed.
	deadline := time.Now().Add(20 * time.Second)
	for time.Now().Before(deadline) {
		if _, err := os.Stat(filepath.Join(dir, "a.gguf")); err == nil {
			break
		}
		time.Sleep(10 * time.Millisecond)
	}
	if err := m.Cancel(id); err != nil {
		t.Fatal(err)
	}
	waitJob(t, m, id)

	if _, err := os.Stat(filepath.Join(dir, "a.gguf")); !os.IsNotExist(err) {
		t.Error("cancel kept a file the job itself downloaded")
	}
	if got, err := os.ReadFile(filepath.Join(dir, "pre.gguf")); err != nil || !bytes.Equal(got, pre) {
		t.Errorf("cancel deleted a file that predated the job: err %v", err)
	}
}

func TestSweepPartials_AgeGated(t *testing.T) {
	root := t.TempDir()
	dir := filepath.Join(root, "o", "r")
	if err := os.MkdirAll(dir, 0o755); err != nil {
		t.Fatal(err)
	}
	old := filepath.Join(dir, "old.gguf"+partSuffix)
	fresh := filepath.Join(dir, "fresh.gguf"+partSuffix)
	keep := filepath.Join(dir, "model.gguf")
	for _, p := range []string{old, fresh, keep} {
		if err := os.WriteFile(p, blobOf(1024), 0o644); err != nil {
			t.Fatal(err)
		}
	}
	stale := time.Now().Add(-48 * time.Hour)
	if err := os.Chtimes(old, stale, stale); err != nil {
		t.Fatal(err)
	}

	n, freed := SweepPartials(root, 24*time.Hour, nil)
	if n != 1 || freed != 1024 {
		t.Errorf("swept %d files / %d bytes, want 1 / 1024", n, freed)
	}
	if _, err := os.Stat(old); !os.IsNotExist(err) {
		t.Error("the stale partial survived the sweep")
	}
	// A partial is resumable, so a recent one is work in progress, not garbage.
	if _, err := os.Stat(fresh); err != nil {
		t.Error("the sweep ate a fresh partial")
	}
	if _, err := os.Stat(keep); err != nil {
		t.Error("the sweep touched a real model file")
	}
}

func TestManager_QueuesSecondJobForSameRepo(t *testing.T) {
	release := make(chan struct{})
	var hits atomic.Int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		<-release
		_, _ = w.Write(make([]byte, 1<<20))
	}))
	defer srv.Close()

	root := t.TempDir()
	src := &fakeSource{base: srv.URL, files: []File{
		{Path: "a.gguf", SizeBytes: 1 << 20},
		{Path: "b.gguf", SizeBytes: 1 << 20},
	}}
	m := NewManager(func() string { return root }, nil, src)

	first, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"a.gguf"}})
	if err != nil {
		t.Fatal(err)
	}
	// Picking a second file from the same repo is the ordinary case (two quants,
	// or a model and its projector) — it waits its turn instead of erroring.
	second, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"b.gguf"}})
	if err != nil {
		t.Fatalf("second job for the same repo was refused: %v", err)
	}
	if j, _ := m.Job(second); j.Phase != PhaseQueued {
		t.Errorf("queued job is in phase %q, want %q", j.Phase, PhaseQueued)
	}
	if got := hits.Load(); got > 1 {
		t.Errorf("%d concurrent transfers for one repo, want 1", got)
	}

	close(release)
	waitJob(t, m, first)
	waitJob(t, m, second)
	if j, _ := m.Job(second); j.Phase != PhaseDone {
		t.Errorf("queued job ended in %q (%s), want %q", j.Phase, j.Err, PhaseDone)
	}
	if _, err := os.Stat(filepath.Join(root, "o", "r", "b.gguf")); err != nil {
		t.Errorf("the queued job never ran: %v", err)
	}
}

// A job canceled before its turn takes nothing with it: the bytes under that
// repo belong to whatever is running there.
func TestManager_CancelQueuedJob(t *testing.T) {
	release := make(chan struct{})
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		<-release
	}))
	defer srv.Close()
	defer close(release)

	root := t.TempDir()
	src := &fakeSource{base: srv.URL, files: []File{
		{Path: "a.gguf", SizeBytes: 1 << 20},
		{Path: "b.gguf", SizeBytes: 1 << 20},
	}}
	m := NewManager(func() string { return root }, nil, src)

	first, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"a.gguf"}})
	if err != nil {
		t.Fatal(err)
	}
	second, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"b.gguf"}})
	if err != nil {
		t.Fatal(err)
	}
	if err := m.Cancel(second); err != nil {
		t.Fatalf("canceling a queued job: %v", err)
	}
	if j, _ := m.Job(second); j.Phase != PhaseCanceled {
		t.Errorf("canceled queued job is in phase %q", j.Phase)
	}
	// The running job is untouched by the other job's cancel.
	if j, _ := m.Job(first); j.Done() {
		t.Errorf("canceling the queued job stopped the running one (phase %q)", j.Phase)
	}
	m.Cancel(first)
	waitJob(t, m, first)
}

// waitForBytes waits until size(path) is positive, so a test can pause a held
// transfer knowing some bytes really landed rather than racing the cancel.
func waitForBytes(t *testing.T, path string) int64 {
	t.Helper()
	deadline := time.Now().Add(20 * time.Second)
	for time.Now().Before(deadline) {
		if n := partSize(path); n > 0 {
			return n
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatalf("no bytes ever landed in %s", path)
	return 0
}

// A download interrupted by a process kill leaves its `.part` and a journal
// record beside it. Restore has to turn that back into a paused job whose
// progress is read off the disk, and the existing Resume verb has to finish it
// — this is the whole point of the journal.
func TestManager_RestoreResumesPartial(t *testing.T) {
	blob := blobOf(1 << 20)
	srv, stop := heldServer(t, blob, 4096)

	root := t.TempDir()
	src := &fakeSource{base: srv.URL, files: []File{{Path: "m.gguf", SizeBytes: int64(len(blob))}}}
	m := NewManager(func() string { return root }, nil, src)

	id, err := m.Start(context.Background(), StartRequest{
		Source: "fake", Repo: "o/r", Files: []string{"m.gguf"}, Label: "m.gguf",
	})
	if err != nil {
		t.Fatal(err)
	}
	part := filepath.Join(root, "o", "r", "m.gguf"+partSuffix)
	waitForBytes(t, part)
	if err := m.Pause(id); err != nil {
		t.Fatal(err)
	}
	waitPhase(t, m, id, PhasePaused)

	// A fresh Manager over the same folder stands in for a restarted process:
	// the old job list is gone, only the disk survived.
	m2 := NewManager(func() string { return root }, nil, src)
	if n := m2.Restore(24 * time.Hour); n != 1 {
		t.Fatalf("Restore returned %d jobs, want 1", n)
	}
	job, ok := m2.Job(id)
	if !ok {
		t.Fatalf("the journaled job %s was not restored", id)
	}
	if job.Phase != PhasePaused {
		t.Errorf("restored job is %q, want paused", job.Phase)
	}
	if job.Label != "m.gguf" || job.Repo != "o/r" || job.Source != "fake" {
		t.Errorf("restored job lost its identity: %+v", job)
	}
	if job.Downloaded <= 0 || job.Downloaded >= job.Total {
		t.Errorf("restored progress = %d/%d, want a partial count", job.Downloaded, job.Total)
	}
	if job.Files[0].Done != job.Downloaded {
		t.Errorf("file progress %d disagrees with job progress %d", job.Files[0].Done, job.Downloaded)
	}

	// Resume against a server that answers in full; the bytes already on disk
	// must carry over rather than being fetched again.
	stop()
	full := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Range") == "" {
			t.Error("resume did not send a Range header")
		}
		serveBlob(w, r, blob, 0)
	}))
	defer full.Close()
	src.base = full.URL

	if err := m2.Resume(id); err != nil {
		t.Fatal(err)
	}
	got := waitJob(t, m2, id)
	if got.Phase != PhaseDone {
		t.Fatalf("resumed job ended %q (%s), want done", got.Phase, got.Err)
	}
	b, err := os.ReadFile(filepath.Join(root, "o", "r", "m.gguf"))
	if err != nil || !bytes.Equal(b, blob) {
		t.Errorf("resumed file mismatch: err %v, %d bytes", err, len(b))
	}
}

// Restore must rebuild per-file progress: a shard that finished before the kill
// counts complete, the one that was running counts its `.part`, and the two add
// up to the row's progress.
func TestManager_RestoreKeepsFinishedFiles(t *testing.T) {
	finished, running := blobOf(2048), blobOf(1<<20)
	release := make(chan struct{})
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasSuffix(r.URL.Path, "a.gguf") {
			serveBlob(w, r, finished, 0)
			return
		}
		w.Write(running[:4096])
		if fl, ok := w.(http.Flusher); ok {
			fl.Flush()
		}
		<-release
	}))
	defer srv.Close()
	defer close(release)
	srv.Config.ErrorLog = discardLogger()

	root := t.TempDir()
	dir := filepath.Join(root, "o", "r")
	src := &fakeSource{base: srv.URL, files: []File{
		{Path: "a.gguf", SizeBytes: int64(len(finished))},
		{Path: "b.gguf", SizeBytes: int64(len(running))},
	}}
	m := NewManager(func() string { return root }, nil, src)
	id, err := m.Start(context.Background(), StartRequest{
		Source: "fake", Repo: "o/r", Files: []string{"a.gguf", "b.gguf"},
	})
	if err != nil {
		t.Fatal(err)
	}
	// Wait for a.gguf to land and b.gguf's partial to start, then stop mid-flight.
	waitForBytes(t, filepath.Join(dir, "b.gguf"+partSuffix))
	deadline := time.Now().Add(20 * time.Second)
	for time.Now().Before(deadline) && existingSize(filepath.Join(dir, "a.gguf")) < int64(len(finished)) {
		time.Sleep(10 * time.Millisecond)
	}
	if existingSize(filepath.Join(dir, "a.gguf")) != int64(len(finished)) {
		t.Fatal("the first shard never landed")
	}
	if err := m.Pause(id); err != nil {
		t.Fatal(err)
	}
	waitPhase(t, m, id, PhasePaused)

	m2 := NewManager(func() string { return root }, nil, src)
	if n := m2.Restore(24 * time.Hour); n != 1 {
		t.Fatalf("Restore returned %d jobs, want 1", n)
	}
	job, _ := m2.Job(id)
	if !job.Files[0].Skipped || job.Files[0].Done != int64(len(finished)) {
		t.Errorf("finished shard not restored as complete: %+v", job.Files[0])
	}
	if job.Files[1].Done <= 0 || job.Files[1].Done >= int64(len(running)) {
		t.Errorf("partial shard progress = %d, want a partial count", job.Files[1].Done)
	}
	if want := int64(len(finished)) + job.Files[1].Done; job.Downloaded != want {
		t.Errorf("job progress = %d, want %d", job.Downloaded, want)
	}
	if job.Total != int64(len(finished)+len(running)) {
		t.Errorf("job total = %d, want %d", job.Total, len(finished)+len(running))
	}

	// The finished shard is removed while the job is parked. Resume must refetch
	// it rather than trust the restored Skipped flag — that flag described the
	// disk at restore time, not now.
	if err := os.Remove(filepath.Join(dir, "a.gguf")); err != nil {
		t.Fatal(err)
	}
	full := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasSuffix(r.URL.Path, "a.gguf") {
			serveBlob(w, r, finished, 0)
			return
		}
		serveBlob(w, r, running, 0)
	}))
	defer full.Close()
	src.base = full.URL
	if err := m2.Resume(id); err != nil {
		t.Fatal(err)
	}
	got := waitJob(t, m2, id)
	if got.Phase != PhaseDone {
		t.Fatalf("resumed job ended %q (%s), want done", got.Phase, got.Err)
	}
	for name, want := range map[string][]byte{"a.gguf": finished, "b.gguf": running} {
		b, err := os.ReadFile(filepath.Join(dir, name))
		if err != nil || !bytes.Equal(b, want) {
			t.Errorf("%s after resume: err %v, %d bytes", name, err, len(b))
		}
	}
}

// Two jobs share one repo folder and therefore one journal file, and both have
// to come back after a restart.
func TestManager_RestoreKeepsBothJobsForOneRepo(t *testing.T) {
	blob := blobOf(1 << 20)
	srv, stop := heldServer(t, blob, 4096)
	defer stop()

	root := t.TempDir()
	src := &fakeSource{base: srv.URL, files: []File{
		{Path: "a.gguf", SizeBytes: int64(len(blob))},
		{Path: "b.gguf", SizeBytes: int64(len(blob))},
	}}
	m := NewManager(func() string { return root }, nil, src)
	first, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"a.gguf"}, Label: "a"})
	if err != nil {
		t.Fatal(err)
	}
	second, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"b.gguf"}, Label: "b"})
	if err != nil {
		t.Fatal(err)
	}
	waitForBytes(t, filepath.Join(root, "o", "r", "a.gguf"+partSuffix))
	if err := m.Pause(first); err != nil {
		t.Fatal(err)
	}
	waitPhase(t, m, first, PhasePaused)
	// The pause frees the repo, so the queued job starts on its own.
	waitForBytes(t, filepath.Join(root, "o", "r", "b.gguf"+partSuffix))
	if err := m.Pause(second); err != nil {
		t.Fatal(err)
	}
	waitPhase(t, m, second, PhasePaused)

	m2 := NewManager(func() string { return root }, nil, src)
	if n := m2.Restore(24 * time.Hour); n != 2 {
		t.Fatalf("Restore returned %d jobs, want 2", n)
	}
	for _, id := range []string{first, second} {
		if _, ok := m2.Job(id); !ok {
			t.Errorf("job %s was not restored", id)
		}
	}
}

// A record whose files are all on disk at their recorded revision is work that
// already finished — the process died between the last rename and the journal
// clear. It must not come back as a paused row.
func TestManager_RestoreDropsFinishedRecord(t *testing.T) {
	root := t.TempDir()
	dir := filepath.Join(root, "o", "r")
	if err := os.MkdirAll(dir, 0o755); err != nil {
		t.Fatal(err)
	}
	blob := blobOf(100)
	if err := os.WriteFile(filepath.Join(dir, "m.gguf"), blob, 0o644); err != nil {
		t.Fatal(err)
	}
	m := NewManager(func() string { return root }, nil)
	m.recordManifest(dir, "m.gguf", manifestEntry{Size: int64(len(blob)), OID: "oid-1"})
	m.recordDownload(dir, downloadRecord{
		ID: "dl-finished", Source: "fake", Repo: "o/r", Started: time.Now(),
		Files: []recordFile{{Path: "m.gguf", Size: int64(len(blob)), OID: "oid-1"}},
	})
	if n := m.Restore(24 * time.Hour); n != 0 {
		t.Errorf("restored %d jobs for a finished download", n)
	}
	if _, err := os.Stat(filepath.Join(dir, journalName)); !os.IsNotExist(err) {
		t.Error("the finished job's journal record survived restore")
	}
}

// A record with no bytes anywhere has nothing to resume: a job killed before
// its first byte, or a `.part` the sweep took. It is cleared, not rendered as
// an empty paused row for the life of the install.
func TestManager_RestoreDropsEmptyRecord(t *testing.T) {
	root := t.TempDir()
	dir := filepath.Join(root, "o", "r")
	m := NewManager(func() string { return root }, nil)
	m.recordDownload(dir, downloadRecord{
		ID: "dl-empty", Source: "fake", Repo: "o/r", Started: time.Now(),
		Files: []recordFile{{Path: "m.gguf", Size: 4096}},
	})
	if n := m.Restore(24 * time.Hour); n != 0 {
		t.Errorf("restored %d jobs with no bytes on disk", n)
	}
	if _, err := os.Stat(filepath.Join(dir, journalName)); !os.IsNotExist(err) {
		t.Error("an empty job's journal record survived restore")
	}
}

// Restore and SweepPartials share the age gate: a `.part` old enough for the
// sweep is not progress, or the row would promise bytes about to be deleted.
func TestManager_RestoreIgnoresAgedPartial(t *testing.T) {
	root := t.TempDir()
	dir := filepath.Join(root, "o", "r")
	if err := os.MkdirAll(dir, 0o755); err != nil {
		t.Fatal(err)
	}
	part := filepath.Join(dir, "m.gguf"+partSuffix)
	if err := os.WriteFile(part, blobOf(100), 0o644); err != nil {
		t.Fatal(err)
	}
	old := time.Now().Add(-48 * time.Hour)
	if err := os.Chtimes(part, old, old); err != nil {
		t.Fatal(err)
	}
	m := NewManager(func() string { return root }, nil)
	m.recordDownload(dir, downloadRecord{
		ID: "dl-old", Source: "fake", Repo: "o/r", Started: time.Now(),
		Files: []recordFile{{Path: "m.gguf", Size: 4096}},
	})
	if n := m.Restore(24 * time.Hour); n != 0 {
		t.Errorf("restored %d job(s) off a partial the sweep will delete", n)
	}
}

// LocalFiles is what puts "downloaded" on a picker row, so it has to agree with
// what the transfer actually leaves behind: a finished file counts, a `.part`
// does not.
func TestManager_LocalFiles(t *testing.T) {
	root := t.TempDir()
	dir := filepath.Join(root, "o", "r", "sub")
	if err := os.MkdirAll(dir, 0o755); err != nil {
		t.Fatal(err)
	}
	for name, size := range map[string]int{
		filepath.Join(root, "o", "r", "done.gguf"):            2048,
		filepath.Join(root, "o", "r", "half.gguf"+partSuffix): 10,
		filepath.Join(dir, "nested.gguf"):                     7,
	} {
		if err := os.WriteFile(name, make([]byte, size), 0o644); err != nil {
			t.Fatal(err)
		}
	}
	m := NewManager(func() string { return root }, nil)
	// The journal is bookkeeping like the manifest: it must not read as a repo
	// file, or the picker would show a row for it.
	m.recordDownload(filepath.Join(root, "o", "r"), downloadRecord{
		ID: "dl-x", Source: "fake", Repo: "o/r", Started: time.Now(),
		Files: []recordFile{{Path: "done.gguf", Size: 2048}},
	})
	got := m.LocalFiles("o/r")
	if got["done.gguf"].Size != 2048 {
		t.Errorf("done.gguf = %d, want 2048", got["done.gguf"].Size)
	}
	if got["sub/nested.gguf"].Size != 7 {
		t.Errorf("nested file missing or wrong size: %v", got)
	}
	if _, ok := got["half.gguf"]; ok {
		t.Error("a .part was reported as a local file")
	}
	if _, ok := got[journalName]; ok {
		t.Error("the download journal was reported as a local file")
	}
	// Nothing recorded these, so they carry no identity — which is what keeps
	// a hand-copied file from ever being called stale.
	if got["done.gguf"].OID != "" {
		t.Errorf("an unrecorded file reported an identity: %q", got["done.gguf"].OID)
	}
	if len(m.LocalFiles("nobody/here")) != 0 {
		t.Error("an unknown repo reported local files")
	}
}

// A manifest is only believed while it still describes the file next to it.
// Recording an id and then letting the bytes change behind our back must drop
// the id, not keep vouching for a file that is gone.
func TestManager_LocalFilesManifest(t *testing.T) {
	root := t.TempDir()
	dir := filepath.Join(root, "o", "r")
	if err := os.MkdirAll(dir, 0o755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(dir, "a.gguf"), make([]byte, 100), 0o644); err != nil {
		t.Fatal(err)
	}
	m := NewManager(func() string { return root }, nil)
	m.recordManifest(dir, "a.gguf", manifestEntry{Size: 100, OID: "sha-one"})

	if got := m.LocalFiles("o/r"); got["a.gguf"].OID != "sha-one" {
		t.Errorf("recorded identity not reported: %+v", got["a.gguf"])
	}
	// The manifest itself is our bookkeeping, not a repo file.
	if _, ok := m.LocalFiles("o/r")[manifestName]; ok {
		t.Error("the manifest was reported as a downloaded file")
	}
	// Same name, different bytes, written by something that is not us.
	if err := os.WriteFile(filepath.Join(dir, "a.gguf"), make([]byte, 140), 0o644); err != nil {
		t.Fatal(err)
	}
	if got := m.LocalFiles("o/r"); got["a.gguf"].OID != "" {
		t.Errorf("a replaced file kept the old identity: %+v", got["a.gguf"])
	}
}

// The bug this whole path exists for: a repo re-uploads a quant under the same
// name, at a size a byte comparison cannot tell apart. Skipping it is what used
// to force the user to rename or delete the file by hand.
func TestManager_HaveCurrent(t *testing.T) {
	root := t.TempDir()
	dir := filepath.Join(root, "o", "r")
	if err := os.MkdirAll(dir, 0o755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(dir, "a.gguf"), make([]byte, 100), 0o644); err != nil {
		t.Fatal(err)
	}
	m := NewManager(func() string { return root }, nil)
	m.recordManifest(dir, "a.gguf", manifestEntry{Size: 100, OID: "old"})
	man := readManifest(dir)

	cases := []struct {
		name  string
		f     JobFile
		force bool
		want  bool
	}{
		{"same revision is skipped", JobFile{Path: "a.gguf", Size: 100, OID: "old"}, false, true},
		{"replaced upstream at the same size is refetched", JobFile{Path: "a.gguf", Size: 100, OID: "new"}, false, false},
		{"force refetches regardless", JobFile{Path: "a.gguf", Size: 100, OID: "old"}, true, false},
		{"a hub stating no id falls back to size", JobFile{Path: "a.gguf", Size: 100}, false, true},
		{"a truncated copy is refetched", JobFile{Path: "a.gguf", Size: 4096, OID: "old"}, false, false},
		{"a file we do not have is fetched", JobFile{Path: "b.gguf", Size: 100, OID: "new"}, false, false},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			if got := m.haveCurrent(dir, tc.f, man, tc.force); got != tc.want {
				t.Errorf("haveCurrent = %v, want %v", got, tc.want)
			}
		})
	}

	// No manifest at all — the hand-copied case. There is nothing to compare
	// against, so the size test stands and the file is left alone.
	if !m.haveCurrent(dir, JobFile{Path: "a.gguf", Size: 100, OID: "new"}, manifest{Files: map[string]manifestEntry{}}, false) {
		t.Error("an unrecorded file was refetched on an id we never had")
	}
}

func TestManager_RejectsFileNotInRepo(t *testing.T) {
	root := t.TempDir()
	src := &fakeSource{base: "http://127.0.0.1:1", files: []File{{Path: "m.gguf", SizeBytes: 10}}}
	m := NewManager(func() string { return root }, nil, src)
	// A path the hub never listed has no size, so neither the disk check nor
	// the progress bar would have a denominator — and it is how a crafted
	// request would try to reach an arbitrary file.
	if _, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"../../etc/passwd"}}); err == nil {
		t.Error("Start accepted a file that is not in the repo listing")
	}
}

func TestManager_NoModelsRootIsAClearError(t *testing.T) {
	src := &fakeSource{base: "http://127.0.0.1:1", files: []File{{Path: "m.gguf", SizeBytes: 10}}}
	m := NewManager(func() string { return "" }, nil, src)
	_, err := m.Start(context.Background(), StartRequest{Source: "fake", Repo: "o/r", Files: []string{"m.gguf"}})
	if err == nil || !strings.Contains(err.Error(), "models folder") {
		t.Errorf("err = %v, want a models-folder message", err)
	}
}

// Clearing the panel's history drops the terminal rows and nothing else: a
// paused pull is outstanding work, and losing it from the list would be the
// same as losing the download.
func TestManager_ClearFinishedKeepsUnfinishedJobs(t *testing.T) {
	m := NewManager(func() string { return t.TempDir() }, nil)
	seed := func(id, phase string) {
		m.jobs[id] = &Job{ID: id, Phase: phase}
		m.order = append(m.order, id)
	}
	seed("a", PhaseDone)
	seed("b", PhaseDownloading)
	seed("c", PhaseCanceled)
	seed("d", PhasePaused)
	seed("e", PhaseError)
	seed("f", PhaseQueued)

	if n := m.ClearFinished(); n != 3 {
		t.Errorf("cleared %d, want the 3 terminal jobs", n)
	}
	var left []string
	for _, j := range m.Jobs() {
		left = append(left, j.ID)
	}
	want := []string{"b", "d", "f"}
	if len(left) != len(want) {
		t.Fatalf("left = %v, want %v", left, want)
	}
	for i, id := range want {
		if left[i] != id {
			t.Fatalf("left = %v, want %v (order preserved)", left, want)
		}
	}
	// A second pass has nothing to do rather than something to break.
	if n := m.ClearFinished(); n != 0 {
		t.Errorf("second clear removed %d, want 0", n)
	}
}