llmgw/internal/mirror/api_test.go

package mirror

import (
	"encoding/json"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

// testAPI is an API over an empty mirror, and the root it holds.
func testAPI(t *testing.T) (*API, string) {
	t.Helper()
	root := t.TempDir()
	m := NewManager(root, discardLog())
	// The upstreams in these tests are httptest servers on loopback.
	m.AllowPrivateSources()
	a := NewAPI(m, "http://mirror:8081")

	// A download outlives the request that started it — that is the point of
	// the manager. So the temp directory must not be torn down while one is
	// still writing into it. Registered after t.TempDir so it runs first.
	t.Cleanup(func() {
		deadline := time.Now().Add(5 * time.Second)
		for time.Now().Before(deadline) {
			all, err := a.m.All()
			if err != nil {
				return
			}
			busy := false
			for _, st := range all {
				busy = busy || st.State == Fetching
			}
			if !busy {
				return
			}
			time.Sleep(5 * time.Millisecond)
		}
		t.Error("a download was still running when the test ended")
	})
	return a, root
}

func call(t *testing.T, a *API, method, path, body string) (int, map[string]any) {
	t.Helper()
	var r io.Reader
	if body != "" {
		r = strings.NewReader(body)
	}
	req := httptest.NewRequest(method, path, r)
	w := httptest.NewRecorder()
	a.Handler().ServeHTTP(w, req)

	var out map[string]any
	if err := json.Unmarshal(w.Body.Bytes(), &out); err != nil {
		t.Fatalf("%s %s returned %d with an unreadable body %q: %v",
			method, path, w.Code, w.Body.String(), err)
	}
	return w.Code, out
}

// waitFor polls until the condition holds, because a download runs in its own
// goroutine and the test would otherwise be asserting on a race.
func waitFor(t *testing.T, what string, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Fatalf("timed out waiting for %s", what)
}

func TestFetchAnswersBeforeTheBytesArrive(t *testing.T) {
	a, root := testAPI(t)
	source := serving(t, "weights")

	code, body := call(t, a, "POST", "/v1/models", `{
		"source": "`+source+`", "model": "llama", "version": "2026-01-01"}`)

	// 202 rather than 200: a caller that reads this as "done" points a catalog
	// at a file that is still arriving.
	if code != http.StatusAccepted {
		t.Fatalf("expected 202 while the download runs, got %d: %v", code, body)
	}
	if got := body["url"]; got != "http://mirror:8081/llama/2026-01-01/model.gguf" {
		t.Errorf("the response should say where the catalog will point; got %v", got)
	}

	waitFor(t, "the download to land on disk", func() bool {
		_, err := os.Stat(filepath.Join(root, "llama", "2026-01-01", "model.gguf"))
		return err == nil
	})
}

func TestTheFilenameComesFromTheSourceWhenNotGiven(t *testing.T) {
	a, _ := testAPI(t)
	_, body := call(t, a, "POST", "/v1/models",
		`{"source": "`+serving(t, "w")+`", "model": "m", "version": "v"}`)

	if got := body["url"]; !strings.HasSuffix(got.(string), "/m/v/model.gguf") {
		t.Errorf("the filename should be taken from the source URL; got %v", got)
	}
}

func TestFetchRefusesWithoutAModelAndVersion(t *testing.T) {
	a, _ := testAPI(t)
	code, body := call(t, a, "POST", "/v1/models", `{"source": "http://x/model.gguf"}`)

	if code != http.StatusBadRequest {
		t.Fatalf("expected 400 without a model and version, got %d", code)
	}
	// The reason belongs in the body: a bare 400 turns a typo into an hour of
	// finding out which field was wrong.
	if !strings.Contains(body["error"].(string), "version") {
		t.Errorf("the error should name what is missing; got %q", body["error"])
	}
}

func TestAskingTwiceDoesNotStartTwoDownloads(t *testing.T) {
	a, _ := testAPI(t)
	release := make(chan struct{})
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		<-release
		_, _ = w.Write([]byte("weights"))
	}))
	defer srv.Close()
	defer close(release)

	body := `{"source": "` + srv.URL + `/model.gguf", "model": "m", "version": "v"}`
	if code, _ := call(t, a, "POST", "/v1/models", body); code != http.StatusAccepted {
		t.Fatalf("the first fetch should be accepted, got %d", code)
	}
	waitFor(t, "the first download to register", func() bool {
		st, ok := a.m.Status("m", "v")
		return ok && st.State == Fetching
	})

	// The second must join the first rather than open a rival writer on the
	// same .part file.
	code, second := call(t, a, "POST", "/v1/models", body)
	if code != http.StatusAccepted {
		t.Fatalf("asking again should report the running download, got %d: %v", code, second)
	}
	st := second["status"].(map[string]any)
	if st["state"] != string(Fetching) {
		t.Errorf("the second answer should describe the download already running; got %v", st["state"])
	}
}

func TestAFailedFetchIsRememberedWithItsReason(t *testing.T) {
	a, _ := testAPI(t)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		http.Error(w, "gone", http.StatusNotFound)
	}))
	defer srv.Close()

	call(t, a, "POST", "/v1/models",
		`{"source": "`+srv.URL+`/model.gguf", "model": "m", "version": "v"}`)

	// Remembered rather than forgotten: a control plane waiting on a model that
	// is never coming must be able to find out why.
	waitFor(t, "the failure to be recorded", func() bool {
		st, ok := a.m.Status("m", "v")
		return ok && st.State == Failed
	})
	code, body := call(t, a, "GET", "/v1/models/m/v", "")
	if code != http.StatusOK {
		t.Fatalf("a failed fetch should still be reportable, got %d", code)
	}
	st := body["status"].(map[string]any)
	if !strings.Contains(st["error"].(string), "404") {
		t.Errorf("the recorded error should say what upstream answered; got %q", st["error"])
	}
}

func TestStatusReadsTheDiskForModelsThisProcessDidNotFetch(t *testing.T) {
	a, root := testAPI(t)
	// A restart forgets every download. A model that is present must not be
	// reported as unknown merely because this process did not put it there.
	writeFile(t, filepath.Join(root, "m", "v", "model.gguf"), "weights")

	code, body := call(t, a, "GET", "/v1/models/m/v", "")
	if code != http.StatusOK {
		t.Fatalf("a model on disk should be found, got %d: %v", code, body)
	}
	if st := body["status"].(map[string]any); st["state"] != string(Present) {
		t.Errorf("expected present, got %v", st["state"])
	}
}

func TestAnUnknownModelIsNotFound(t *testing.T) {
	a, _ := testAPI(t)
	if code, _ := call(t, a, "GET", "/v1/models/nope/v1", ""); code != http.StatusNotFound {
		t.Errorf("expected 404 for a model the mirror does not hold, got %d", code)
	}
}

func TestListReportsTheBaseSoCatalogsAgree(t *testing.T) {
	a, root := testAPI(t)
	writeFile(t, filepath.Join(root, "m", "v", "model.gguf"), "weights")

	code, body := call(t, a, "GET", "/v1/models", "")
	if code != http.StatusOK {
		t.Fatalf("expected 200, got %d", code)
	}
	// Returned rather than assumed: a caller reconstructing the URL with a
	// different spelling is what makes a file look unused when it is not.
	if body["base"] != "http://mirror:8081" {
		t.Errorf("the listing should say where nodes fetch from; got %v", body["base"])
	}
	if n := len(body["models"].([]any)); n != 1 {
		t.Errorf("expected the one model on disk, got %d", n)
	}
}

func TestRemoveTakesEveryFileOfAVersion(t *testing.T) {
	a, root := testAPI(t)
	// A model split across files is useless with one of them missing, so the
	// version is the unit of removal, not the file.
	writeFile(t, filepath.Join(root, "m", "v", "part-1.safetensors"), "a")
	writeFile(t, filepath.Join(root, "m", "v", "part-2.safetensors"), "b")

	code, body := call(t, a, "DELETE", "/v1/models/m/v", "")
	if code != http.StatusOK {
		t.Fatalf("expected 200, got %d: %v", code, body)
	}
	if n := len(body["removed"].([]any)); n != 2 {
		t.Errorf("both files should have gone; removed %v", body["removed"])
	}
	if _, err := os.Stat(filepath.Join(root, "m")); !os.IsNotExist(err) {
		t.Error("the emptied model directory should have gone with it")
	}
}

func TestRemoveRefusesWhileADownloadIsRunning(t *testing.T) {
	a, _ := testAPI(t)
	release := make(chan struct{})
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		<-release
		_, _ = w.Write([]byte("weights"))
	}))
	defer srv.Close()
	defer close(release)

	call(t, a, "POST", "/v1/models",
		`{"source": "`+srv.URL+`/model.gguf", "model": "m", "version": "v"}`)
	waitFor(t, "the download to register", func() bool {
		st, ok := a.m.Status("m", "v")
		return ok && st.State == Fetching
	})

	// Deleting under a running fetch lets the file reappear when the download
	// finishes, which looks like the delete silently failed.
	code, body := call(t, a, "DELETE", "/v1/models/m/v", "")
	if code != http.StatusConflict {
		t.Fatalf("expected 409 while downloading, got %d: %v", code, body)
	}
}

func TestRemovingSomethingAbsentIsNotFound(t *testing.T) {
	a, _ := testAPI(t)
	if code, _ := call(t, a, "DELETE", "/v1/models/m/v", ""); code != http.StatusNotFound {
		t.Errorf("expected 404, got %d", code)
	}
}

func TestHealthFailsWhenTheMirrorDirectoryIsGone(t *testing.T) {
	a, root := testAPI(t)
	if code, _ := call(t, a, "GET", "/v1/health", ""); code != http.StatusOK {
		t.Fatalf("a readable mirror should be healthy, got %d", code)
	}

	// A mirror whose disk has gone away still accepts connections. A check that
	// only proved the process was running would call that healthy right up
	// until a node tried to fetch.
	if err := os.RemoveAll(root); err != nil {
		t.Fatal(err)
	}
	if code, _ := call(t, a, "GET", "/v1/health", ""); code != http.StatusServiceUnavailable {
		t.Errorf("expected 503 once the directory is unreadable, got %d", code)
	}
}

func writeFile(t *testing.T, path, body string) {
	t.Helper()
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(path, []byte(body), 0o600); err != nil {
		t.Fatal(err)
	}
}

// discardLog writes nothing, so a test that provokes a failure on purpose does
// not print an error line and make a passing run look broken.
func discardLog() *slog.Logger {
	return slog.New(slog.NewTextHandler(io.Discard, nil))
}

func TestTheChecksumSurvivesTheDownloadFinishing(t *testing.T) {
	a, _ := testAPI(t)
	call(t, a, "POST", "/v1/models",
		`{"source": "`+serving(t, "weights")+`", "model": "m", "version": "v"}`)

	waitFor(t, "the download to finish", func() bool {
		st, ok := a.m.Status("m", "v")
		return ok && st.State == Present
	})

	// The listing used to prefer the disk entry wholesale, and a directory has
	// no checksum in it — so the sha256 vanished at the exact moment the file
	// became verifiable. It is the one field an operator compares against what
	// upstream published.
	_, body := call(t, a, "GET", "/v1/models", "")
	models := body["models"].([]any)
	if len(models) != 1 {
		t.Fatalf("expected the one model, got %v", models)
	}
	got := models[0].(map[string]any)
	if got["sha256"] == nil || got["sha256"] == "" {
		t.Errorf("the checksum should still be reported once the file is present; got %v", got)
	}
}

func TestTwoShardsOfOneVersionDownloadTogether(t *testing.T) {
	a, _ := testAPI(t)
	release := make(chan struct{})
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		<-release
		_, _ = w.Write([]byte("weights"))
	}))
	defer srv.Close()
	defer close(release)

	for _, shard := range []string{"part-1.safetensors", "part-2.safetensors"} {
		code, body := call(t, a, "POST", "/v1/models", `{
			"source": "`+srv.URL+`/`+shard+`", "model": "llama", "version": "v1",
			"file": "`+shard+`"}`)
		// Keyed by version, the second shard was refused as already fetching
		// until the first had finished — so a fifteen-shard model had to be
		// pulled one shard at a time, each waiting on the last.
		if code != http.StatusAccepted {
			t.Fatalf("%s was refused with %d: %v", shard, code, body)
		}
		st := body["status"].(map[string]any)
		if st["file"] != shard {
			t.Errorf("expected %s to be its own download, got %v", shard, st["file"])
		}
	}
}

func TestAskingTwiceForTheSameFileStillJoinsTheOneRunning(t *testing.T) {
	a, _ := testAPI(t)
	release := make(chan struct{})
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		<-release
		_, _ = w.Write([]byte("weights"))
	}))
	defer srv.Close()
	defer close(release)

	body := `{"source": "` + srv.URL + `/f.gguf", "model": "m", "version": "v", "file": "f.gguf"}`
	call(t, a, "POST", "/v1/models", body)
	waitFor(t, "the first to register", func() bool {
		st, ok := a.m.Status("m", "v")
		return ok && st.State == Fetching
	})

	// Per file is still single-flight: two writers on one .part file is the
	// thing the key exists to prevent.
	code, second := call(t, a, "POST", "/v1/models", body)
	if code != http.StatusAccepted {
		t.Fatalf("expected the running download to be reported, got %d", code)
	}
	if st := second["status"].(map[string]any); st["state"] != string(Fetching) {
		t.Errorf("expected the one already running, got %v", st["state"])
	}
}

func TestAVersionIsOnlyPresentOnceEveryShardIs(t *testing.T) {
	a, root := testAPI(t)
	writeFile(t, filepath.Join(root, "llama", "v1", "part-1.safetensors"), "a")

	release := make(chan struct{})
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		<-release
		_, _ = w.Write([]byte("b"))
	}))
	defer srv.Close()
	defer close(release)

	call(t, a, "POST", "/v1/models", `{"source": "`+srv.URL+`/part-2.safetensors",
		"model": "llama", "version": "v1", "file": "part-2.safetensors"}`)
	waitFor(t, "the second shard to register", func() bool {
		st, _ := a.m.Status("llama", "v1")
		return st.State == Fetching
	})

	// One shard on disk and one still arriving is not a version anybody can
	// load, so the version reports the worse of the two.
	_, body := call(t, a, "GET", "/v1/models/llama/v1", "")
	if st := body["status"].(map[string]any); st["state"] != string(Fetching) {
		t.Errorf("expected the version to read as fetching, got %v", st["state"])
	}
}

func TestRemovingAVersionIsRefusedWhileAnyShardIsArriving(t *testing.T) {
	a, root := testAPI(t)
	writeFile(t, filepath.Join(root, "llama", "v1", "part-1.safetensors"), "a")

	release := make(chan struct{})
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		<-release
		_, _ = w.Write([]byte("b"))
	}))
	defer srv.Close()
	defer close(release)

	call(t, a, "POST", "/v1/models", `{"source": "`+srv.URL+`/part-2.safetensors",
		"model": "llama", "version": "v1", "file": "part-2.safetensors"}`)
	waitFor(t, "the shard to register", func() bool {
		st, _ := a.m.Status("llama", "v1")
		return st.State == Fetching
	})

	// Deleting the shard that already landed while its sibling is still coming
	// leaves a directory that reappears half-populated.
	if code, _ := call(t, a, "DELETE", "/v1/models/llama/v1", ""); code != http.StatusConflict {
		t.Errorf("expected 409 while a shard is arriving, got %d", code)
	}
}

func TestTheAPIWillNotFetchFromThePrivateNetwork(t *testing.T) {
	// The file server hands whatever was fetched to anyone, so a fetch that
	// could reach 169.254.169.254 or a database on the LAN was a way to read
	// them through the mirror.
	m := NewManager(t.TempDir(), discardLog())
	for _, src := range []string{
		"http://127.0.0.1:1/x.gguf", "http://169.254.169.254/latest/meta-data",
		"http://10.0.0.140:9090/api/v1/nodes", "http://[::1]:1/x",
	} {
		st, err := m.Start(src, Entry{Model: "m", Version: "v", File: "x.gguf"})
		if err != nil {
			t.Fatalf("%s: %v", src, err)
		}
		deadline := time.Now().Add(5 * time.Second)
		for st.State == Fetching && time.Now().Before(deadline) {
			time.Sleep(10 * time.Millisecond)
			st, _ = m.Status("m", "v")
		}
		if st.State != Failed || !strings.Contains(st.Err, "refusing to fetch") {
			t.Errorf("%s: %+v", src, st)
		}
		m.Forget("m", "v")
	}
	for _, src := range []string{"file:///etc/passwd", "gopher://x/", "not a url"} {
		if _, err := m.Start(src, Entry{Model: "m", Version: "v", File: "x"}); err == nil {
			t.Errorf("%s was accepted as a source", src)
		}
	}
}