Loki/internal/loki/backend_models_dl_test.go

package loki

import (
	"bytes"
	"context"
	"math/rand"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"testing"
	"time"
)

// serveBlob serves data with (or without) byte-range support.
func serveBlob(data []byte, ranges bool) *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !ranges {
			w.Header().Set("Content-Length", strconv.Itoa(len(data)))
			_, _ = w.Write(data)
			return
		}
		http.ServeContent(w, r, "m.gguf", time.Time{}, bytes.NewReader(data))
	}))
}

// Un téléchargement annulé ne doit laisser NI le .gguf final NI le .part.
func TestRunDownloadCancelLeavesNothing(t *testing.T) {
	t.Setenv("LOKI_DL_CONNS", "4")
	release := make(chan struct{})
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Range") == "bytes=0-0" { // sonde : réponse immédiate
			w.Header().Set("Content-Range", "bytes 0-0/"+strconv.Itoa(64<<20))
			w.WriteHeader(206)
			_, _ = w.Write([]byte{0})
			return
		}
		// Corps qui traîne : l'annulation doit l'interrompre.
		w.WriteHeader(206)
		w.(http.Flusher).Flush()
		select {
		case <-release:
		case <-r.Context().Done():
		}
	}))
	defer srv.Close()
	defer close(release)

	dir := t.TempDir()
	dest := filepath.Join(dir, "m.gguf")
	ctx, cancel := context.WithCancel(context.Background())
	st := &dlState{Filename: "m.gguf", cancel: cancel}
	done := make(chan struct{})
	go func() { runDownloadSet(ctx, st, []string{srv.URL + "/m.gguf"}, []string{dest}); close(done) }()

	time.Sleep(300 * time.Millisecond)
	cancel()
	select {
	case <-done:
	case <-time.After(10 * time.Second):
		t.Fatal("runDownloadSet n a pas rendu la main apres annulation")
	}

	if !st.Canceled || !st.Finished {
		t.Fatalf("état attendu annulé+terminé, got canceled=%v finished=%v err=%q", st.Canceled, st.Finished, st.Err)
	}
	if _, err := os.Stat(dest + ".part"); !os.IsNotExist(err) {
		t.Fatal(".part laissé sur le disque après annulation")
	}
	if _, err := os.Stat(dest); !os.IsNotExist(err) {
		t.Fatal("fichier final créé alors que le téléchargement a été annulé")
	}
}

// cleanStalePartFiles doit balayer les .part orphelins sans toucher aux .gguf.
func TestCleanStalePartFiles(t *testing.T) {
	testHome(t)
	dir := modelsDir()
	if err := os.MkdirAll(dir, 0o755); err != nil {
		t.Fatal(err)
	}
	keep := filepath.Join(dir, "bon.gguf")
	stale := filepath.Join(dir, "coupe.gguf.part")
	for _, p := range []string{keep, stale} {
		if err := os.WriteFile(p, []byte("x"), 0o644); err != nil {
			t.Fatal(err)
		}
	}
	cleanStalePartFiles()
	if _, err := os.Stat(stale); !os.IsNotExist(err) {
		t.Fatal(".part orphelin non supprimé")
	}
	if _, err := os.Stat(keep); err != nil {
		t.Fatal(".gguf valide supprimé par erreur")
	}
}

func TestRunDownloadParallelAndFallback(t *testing.T) {
	data := make([]byte, 48<<20) // > 3×dlMinChunk so the split actually kicks in
	rand.New(rand.NewSource(1)).Read(data)
	t.Setenv("LOKI_DL_CONNS", "4")

	for _, ranges := range []bool{true, false} {
		srv := serveBlob(data, ranges)
		dir := t.TempDir()
		dest := filepath.Join(dir, "m.gguf")
		st := &dlState{Filename: "m.gguf"}
		runDownloadSet(context.Background(), st, []string{srv.URL + "/m.gguf"}, []string{dest})
		srv.Close()

		if st.Err != "" {
			t.Fatalf("ranges=%v: erreur %s", ranges, st.Err)
		}
		got, err := os.ReadFile(dest)
		if err != nil {
			t.Fatalf("ranges=%v: %v", ranges, err)
		}
		if len(got) != len(data) {
			t.Fatalf("ranges=%v: taille %d != %d", ranges, len(got), len(data))
		}
		for i := range got {
			if got[i] != data[i] {
				t.Fatalf("ranges=%v: octet %d différent", ranges, i)
			}
		}
		if st.Done != int64(len(data)) {
			t.Fatalf("ranges=%v: done=%d", ranges, st.Done)
		}
		if ranges && st.Conns < 2 {
			t.Fatalf("attendu du parallélisme, conns=%d", st.Conns)
		}
		if _, err := os.Stat(dest + ".part"); !os.IsNotExist(err) {
			t.Fatalf("ranges=%v: .part laissé derrière", ranges)
		}
	}
}

// Un dépôt Hugging Face verrouillé répond 401 sur le .gguf alors que son
// arborescence se lit sans jeton. « HTTP 401 depuis la source » n'apprenait
// rien : le message doit nommer le verrou, le dépôt, et le geste à faire.
func TestDLSourceErrorExplainsGatedRepo(t *testing.T) {
	const gated = "https://huggingface.co/orcarouter/Qwen3.8-27B-Uncensored-GGUF/resolve/main/m.gguf"
	resp := func(status int, code string) *http.Response {
		r := &http.Response{StatusCode: status, Header: http.Header{}}
		if code != "" {
			r.Header.Set("X-Error-Code", code)
		}
		return r
	}
	for _, c := range []struct {
		name   string
		token  string
		resp   *http.Response
		url    string
		want   []string
		absent []string
	}{
		{
			name: "gated sans jeton", resp: resp(401, "GatedRepo"), url: gated,
			want: []string{"accès restreint", "orcarouter/Qwen3.8-27B-Uncensored-GGUF", "HF_TOKEN"},
		},
		{
			name: "gated avec jeton refusé", token: "hf_xxx", resp: resp(401, "GatedRepo"), url: gated,
			want:   []string{"accès restreint", "n'y donne pas accès"},
			absent: []string{"renseigne le jeton"},
		},
		{
			name: "dépôt privé ou absent", resp: resp(401, "RepoNotFound"), url: gated,
			want: []string{"privé ou inexistant", "HF_TOKEN"},
		},
		{
			name: "fichier absent", resp: resp(404, "EntryNotFound"), url: gated,
			want: []string{"« m.gguf » absent du dépôt", "404"},
		},
		{
			// Dépôt qui publie six fragments sur sept (table PLE livrée à part) :
			// le message doit NOMMER le fragment manquant, sinon on cherche une
			// révision réécrite qui n'existe pas.
			name: "fragment absent", resp: resp(404, "EntryNotFound"),
			url:  "https://huggingface.co/x/y-GGUF/resolve/main/Q2_K/y-Q2_K-00007-of-00007.gguf",
			want: []string{"« y-Q2_K-00007-of-00007.gguf » absent du dépôt", "fragment jamais publié"},
		},
		{
			name: "source hors Hugging Face", resp: resp(403, ""), url: "https://example.com/m.gguf",
			want:   []string{"accès refusé par la source"},
			absent: []string{"huggingface.co"},
		},
		{
			name: "panne de la source", resp: resp(503, ""), url: gated,
			want: []string{"en panne", "503"},
		},
		{
			name: "code inconnu", resp: resp(418, ""), url: gated,
			want: []string{"HTTP 418 depuis la source"},
		},
	} {
		t.Run(c.name, func(t *testing.T) {
			t.Setenv("HF_TOKEN", c.token)
			got := dlSourceError(c.resp, c.url).Error()
			for _, w := range c.want {
				if !strings.Contains(got, w) {
					t.Errorf("message %q ne contient pas %q", got, w)
				}
			}
			for _, a := range c.absent {
				if strings.Contains(got, a) {
					t.Errorf("message %q contient %q alors qu'il ne devrait pas", got, a)
				}
			}
		})
	}
}

// La sonde est le premier appel réseau d'une installation : c'est elle qui doit
// remonter le refus expliqué, pas un code HTTP nu.
func TestDLProbeSurfacesExplainedError(t *testing.T) {
	t.Setenv("HF_TOKEN", "")
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("X-Error-Code", "GatedRepo")
		w.WriteHeader(401)
	}))
	defer srv.Close()

	_, _, err := dlProbe(context.Background(), srv.URL+"/m.gguf")
	if err == nil {
		t.Fatal("un 401 doit faire échouer la sonde")
	}
	if !strings.Contains(err.Error(), "HF_TOKEN") {
		t.Errorf("message %q sans indication sur le jeton", err)
	}
}

type slowWriter struct{ http.ResponseWriter }

func (s slowWriter) Write(b []byte) (int, error) {
	time.Sleep(2 * time.Millisecond)
	return s.ResponseWriter.Write(b)
}

// Une pause coupe le flux, garde ce qui est écrit, et la reprise termine un
// fichier identique à la source (repris d'AJEAN v0.18.1).
func TestRunDownloadPauseResume(t *testing.T) {
	data := make([]byte, 48<<20)
	rand.New(rand.NewSource(2)).Read(data)
	t.Setenv("LOKI_DL_CONNS", "4")
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.ServeContent(slowWriter{w}, r, "m.gguf", time.Time{}, bytes.NewReader(data))
	}))
	defer srv.Close()

	dest := filepath.Join(t.TempDir(), "m.gguf")
	st := &dlState{Filename: "m.gguf", gate: newDlGate()}
	done := make(chan struct{})
	go func() {
		runDownloadSet(context.Background(), st, []string{srv.URL + "/m.gguf"}, []string{dest})
		close(done)
	}()

	time.Sleep(300 * time.Millisecond)
	st.gate.setPaused(true)
	time.Sleep(1500 * time.Millisecond) // le compteur publié se fige
	dlMu.Lock()
	a := st.Done
	dlMu.Unlock()
	time.Sleep(1500 * time.Millisecond)
	dlMu.Lock()
	b, fin := st.Done, st.Finished
	dlMu.Unlock()
	if fin || a != b || a == 0 {
		t.Fatalf("pause inopérante : %d puis %d octets, terminé=%v", a, b, fin)
	}
	st.gate.setPaused(false)
	select {
	case <-done:
	case <-time.After(30 * time.Second):
		t.Fatal("pas terminé après la reprise")
	}
	if st.Err != "" {
		t.Fatalf("erreur après reprise : %s", st.Err)
	}
	got, err := os.ReadFile(dest)
	if err != nil || !bytes.Equal(got, data) {
		t.Fatalf("fichier différent de la source après reprise (err=%v)", err)
	}
}

// La route : pause puis reprise d'un téléchargement connu, refus pour
// l'inconnu et pour un téléchargement terminé.
func TestHandleModelDownloadPause(t *testing.T) {
	dlMu.Lock()
	saved := dlDownloads
	dlDownloads = map[string]*dlState{
		"a.gguf": {Filename: "a.gguf", Speed: 9, gate: newDlGate()},
		"b.gguf": {Filename: "b.gguf", Finished: true},
	}
	dlMu.Unlock()
	t.Cleanup(func() { dlMu.Lock(); dlDownloads = saved; dlMu.Unlock() })
	post := func(body string) (int, string) {
		w := httptest.NewRecorder()
		handleModelDownloadPause(w, httptest.NewRequest("POST", "/api/models/download/pause", strings.NewReader(body)))
		return w.Code, strings.TrimSpace(w.Body.String())
	}
	if code, out := post(`{"filename":"../a.gguf","pause":true}`); code != 200 || out != `{"ok":true,"paused":true}` {
		t.Fatalf("pause : %d %s", code, out)
	}
	dlMu.Lock()
	st := dlDownloads["a.gguf"]
	paused, speed := st.Paused, st.Speed
	dlMu.Unlock()
	if !paused || speed != 0 || !st.gate.isPaused() {
		t.Errorf("état après pause : paused=%v speed=%d porte=%v", paused, speed, st.gate.isPaused())
	}
	if code, out := post(`{"filename":"a.gguf","pause":false}`); code != 200 || out != `{"ok":true,"paused":false}` || st.gate.isPaused() {
		t.Fatalf("reprise : %d %s", code, out)
	}
	for _, body := range []string{`{"filename":"b.gguf","pause":true}`, `{"filename":"x.gguf","pause":true}`} {
		if code, _ := post(body); code != 404 {
			t.Errorf("%s : %d, attendu 404", body, code)
		}
	}
	if code, _ := post(`[`); code != 400 {
		t.Errorf("corps illisible : %d", code)
	}
}