llmgw/internal/mirror/peer_test.go

package mirror

import (
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
)

// testMirror is a mirror with its file server and management API, as a peer
// reaches it.
type testMirror struct {
	root  string
	m     *Manager
	files *httptest.Server
	api   *httptest.Server
}

func newTestMirror(t *testing.T) *testMirror {
	t.Helper()
	root := t.TempDir()
	tm := &testMirror{root: root, m: NewManager(root, discardLog())}
	tm.m.AllowPrivateSources()
	tm.files = httptest.NewServer(http.FileServer(http.Dir(root)))
	tm.api = httptest.NewServer(NewAPI(tm.m, tm.files.URL).Handler())
	t.Cleanup(func() { tm.files.Close(); tm.api.Close() })
	return tm
}

func (tm *testMirror) client(t *testing.T) *Client {
	c, err := NewClient(tm.api.URL, tlsconf.Config{}, 5*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	return c
}

func (tm *testMirror) has(model, version string) bool {
	st, ok := tm.m.Status(model, version)
	return ok && st.State == Present
}

func waitPresent(t *testing.T, tm *testMirror, model, version string) {
	t.Helper()
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		if tm.has(model, version) {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	st, _ := tm.m.Status(model, version)
	t.Fatalf("%s/%s never arrived: %+v", model, version, st)
}

func TestTwoMirrorsKeepInStepAndAgreeOnWhatWasRemoved(t *testing.T) {
	blob := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write([]byte("weights of a small model"))
	}))
	defer blob.Close()
	a, b := newTestMirror(t), newTestMirror(t)
	if _, err := a.m.Start(blob.URL+"/m.gguf", Entry{Model: "qwen", Version: "v1", File: "m.gguf"}); err != nil {
		t.Fatal(err)
	}
	waitPresent(t, a, "qwen", "v1")

	// b copies what a has, its checksum checked.
	pb := &Peers{M: b.m, Log: discardLog()}
	pb.SyncOnce(t.Context(), a.client(t))
	waitPresent(t, b, "qwen", "v1")
	got, _ := os.ReadFile(filepath.Join(b.root, "qwen", "v1", "m.gguf"))
	if string(got) != "weights of a small model" {
		t.Errorf("b holds %q", got)
	}

	// Removed on a: a does not copy it back from b, and b removes it too.
	time.Sleep(10 * time.Millisecond)
	if _, err := a.m.RemoveVersion("qwen", "v1"); err != nil {
		t.Fatal(err)
	}
	pa := &Peers{M: a.m, Log: discardLog()}
	pa.SyncOnce(t.Context(), b.client(t))
	time.Sleep(200 * time.Millisecond)
	if a.has("qwen", "v1") {
		t.Fatal("a copied back from b what it had removed")
	}
	pb.SyncOnce(t.Context(), a.client(t))
	if b.has("qwen", "v1") {
		t.Error("b kept what a removed")
	}

	// Added again on a, after the removal: the new one stands, on both.
	time.Sleep(10 * time.Millisecond)
	if _, err := a.m.Start(blob.URL+"/m.gguf", Entry{Model: "qwen", Version: "v1", File: "m.gguf"}); err != nil {
		t.Fatal(err)
	}
	waitPresent(t, a, "qwen", "v1")
	pb.SyncOnce(t.Context(), a.client(t))
	waitPresent(t, b, "qwen", "v1")
}

// A file that is not what the peer says it is, is not kept.
func TestACopyWhoseChecksumIsNotThePeersIsNotKept(t *testing.T) {
	served := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write([]byte("something else"))
	}))
	defer served.Close()
	b := newTestMirror(t)
	e := Entry{Model: "qwen", Version: "v1", File: "m.gguf"}
	if _, err := b.m.startCopy(served.URL+"/m.gguf", e, "0000"); err != nil {
		t.Fatal(err)
	}
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		if st, _ := b.m.Status("qwen", "v1"); st.State == Failed {
			break
		}
		time.Sleep(20 * time.Millisecond)
	}
	if _, err := os.Stat(e.Path(b.root)); !os.IsNotExist(err) {
		t.Errorf("a copy with the wrong checksum was kept: %v", err)
	}
}

// A mirror that cannot be reached is passed over for one kept in step with it.
func TestTheNextMirrorIsAskedWhenTheFirstCannotBeReached(t *testing.T) {
	a, b := newTestMirror(t), newTestMirror(t)
	c, err := NewFailoverClient([]string{a.api.URL, b.api.URL}, tlsconf.Config{}, 2*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	if list, err := c.List(t.Context()); err != nil || list.Base != a.files.URL {
		t.Fatalf("both up: %+v %v", list, err)
	}
	a.api.Close()
	if list, err := c.List(t.Context()); err != nil || list.Base != b.files.URL {
		t.Errorf("the first down: %+v %v", list, err)
	}
	b.api.Close()
	if _, err := c.List(t.Context()); err == nil {
		t.Error("both down, and an answer")
	}
}

// A node is given a model's file on a mirror that answers.
func TestAModelsAddressIsMovedToAMirrorThatAnswers(t *testing.T) {
	a, b := newTestMirror(t), newTestMirror(t)
	s, err := NewSources([]string{a.api.URL, b.api.URL}, tlsconf.Config{}, discardLog())
	if err != nil {
		t.Fatal(err)
	}
	s.Check(t.Context())
	src := a.files.URL + "/qwen/v1/m.gguf"
	if got := s.Rewrite(src); got != src {
		t.Errorf("its own mirror up: %s", got)
	}
	a.api.Close()
	s.Check(t.Context())
	if got := s.Rewrite(src); got != b.files.URL+"/qwen/v1/m.gguf" {
		t.Errorf("its own mirror down: %s", got)
	}
	if got := s.Rewrite("https://huggingface.co/x/y"); got != "https://huggingface.co/x/y" {
		t.Errorf("an address on no mirror: %s", got)
	}
	b.api.Close()
	s.Check(t.Context())
	if got := s.Rewrite(src); got != src {
		t.Errorf("none up: %s", got)
	}
}

// A file fetched before checksums were kept is given one, and is then copied.
func TestAFileWithoutAChecksumIsGivenOneAndThenCopied(t *testing.T) {
	a, b := newTestMirror(t), newTestMirror(t)
	path := filepath.Join(a.root, "old", "v1", "m.gguf")
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(path, []byte("fetched long ago"), 0o644); err != nil {
		t.Fatal(err)
	}
	pb := &Peers{M: b.m, Log: discardLog()}
	pb.SyncOnce(t.Context(), a.client(t))
	time.Sleep(100 * time.Millisecond)
	if b.has("old", "v1") {
		t.Fatal("copied without a checksum to check it by")
	}
	if n := a.m.FillChecksums(t.Context()); n != 1 || readChecksum(path) == "" {
		t.Fatalf("recorded %d", n)
	}
	if a.m.FillChecksums(t.Context()) != 0 {
		t.Error("recorded again")
	}
	pb.SyncOnce(t.Context(), a.client(t))
	waitPresent(t, b, "old", "v1")
}

// A file removed outside the running mirror is not reported as here, and its
// removal, recorded, reaches a peer instead of the peer's copy coming back.
func TestAFileRemovedByHandIsNotReportedAndItsRemovalIsCarried(t *testing.T) {
	blob := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write([]byte("weights"))
	}))
	defer blob.Close()
	a, b := newTestMirror(t), newTestMirror(t)
	e := Entry{Model: "qwen", Version: "v1", File: "m.gguf"}
	if _, err := a.m.Start(blob.URL+"/m.gguf", e); err != nil {
		t.Fatal(err)
	}
	waitPresent(t, a, "qwen", "v1")
	pb := &Peers{M: b.m, Log: discardLog()}
	pb.SyncOnce(t.Context(), a.client(t))
	waitPresent(t, b, "qwen", "v1")

	time.Sleep(10 * time.Millisecond)
	if err := Remove(a.root, e); err != nil {
		t.Fatal(err)
	}
	if err := RecordRemoval(a.root, "qwen", "v1"); err != nil {
		t.Fatal(err)
	}
	if a.has("qwen", "v1") {
		t.Fatal("reported as here after it was removed")
	}
	all, _ := a.m.All()
	for _, st := range all {
		if st.Model == "qwen" {
			t.Fatalf("listed after it was removed: %+v", st)
		}
	}
	pa := &Peers{M: a.m, Log: discardLog()}
	pa.SyncOnce(t.Context(), b.client(t))
	time.Sleep(200 * time.Millisecond)
	if a.has("qwen", "v1") {
		t.Fatal("copied back from the peer")
	}
	pb.SyncOnce(t.Context(), a.client(t))
	if b.has("qwen", "v1") {
		t.Error("the peer kept it")
	}
}

// A fetch given the registry's sha256 keeps only that file.
func TestAFetchThatIsNotWhatTheRegistryPublishedIsDeleted(t *testing.T) {
	served := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write([]byte("replaced upstream"))
	}))
	defer served.Close()
	tm := newTestMirror(t)
	e := Entry{Model: "qwen", Version: "v1", File: "m.gguf"}
	if _, err := tm.m.StartChecked(served.URL+"/m.gguf", e, "aaaa"); err != nil {
		t.Fatal(err)
	}
	deadline := time.Now().Add(5 * time.Second)
	var st Status
	for time.Now().Before(deadline) {
		if st, _ = tm.m.Status("qwen", "v1"); st.State == Failed {
			break
		}
		time.Sleep(20 * time.Millisecond)
	}
	if st.State != Failed || !strings.Contains(st.Err, "aaaa") {
		t.Fatalf("%+v", st)
	}
	if _, err := os.Stat(e.Path(tm.root)); !os.IsNotExist(err) {
		t.Error("kept")
	}
}