llmgw/internal/cli/mirror_import_test.go

package cli

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"testing"

	tea "github.com/charmbracelet/bubbletea"
)

// runMirror runs `llmgw mirror ...` as an operator would and returns what it
// printed.
func runMirror(t *testing.T, args ...string) (string, error) {
	t.Helper()
	c := Root()
	var out bytes.Buffer
	c.SetOut(&out)
	c.SetErr(&out)
	c.SetArgs(append([]string{"mirror"}, args...))
	err := c.ExecuteContext(t.Context())
	return out.String(), err
}

// ollamaWith lays out an Ollama models directory holding one model, as
// Ollama writes it, and returns the directory and the weights' sha256.
func ollamaWith(t *testing.T, name, tag string, weights []byte, params string) (string, string) {
	t.Helper()
	dir := filepath.Join(t.TempDir(), "models")
	blob := func(content []byte) string {
		s := sha256.Sum256(content)
		h := hex.EncodeToString(s[:])
		if err := os.MkdirAll(filepath.Join(dir, "blobs"), 0o755); err != nil {
			t.Fatal(err)
		}
		if err := os.WriteFile(filepath.Join(dir, "blobs", "sha256-"+h), content, 0o644); err != nil {
			t.Fatal(err)
		}
		return h
	}
	w := blob(weights)
	p := blob([]byte(params))
	manifest := `{"schemaVersion":2,"mediaType":"application/vnd.docker.distribution.manifest.v2+json",` +
		`"config":{"mediaType":"application/vnd.docker.container.image.v1+json","digest":"sha256:` + blob([]byte("{}")) + `","size":2},` +
		`"layers":[{"mediaType":"application/vnd.ollama.image.model","digest":"sha256:` + w + `","size":` + strconv.Itoa(len(weights)) + `},` +
		`{"mediaType":"application/vnd.ollama.image.params","digest":"sha256:` + p + `","size":` + strconv.Itoa(len(params)) + `}]}`
	mdir := filepath.Join(dir, "manifests", "registry.ollama.ai", "library", name)
	if err := os.MkdirAll(mdir, 0o755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(mdir, tag), []byte(manifest), 0o644); err != nil {
		t.Fatal(err)
	}
	return dir, w
}

func TestImportOllamaListsThenImportsAModelForLlamaCpp(t *testing.T) {
	weights := []byte("GGUF\x03\x00\x00\x00qwen weights")
	ollama, sum := ollamaWith(t, "qwen2.5", "7b", weights, `{"num_ctx":16384}`)
	root := t.TempDir()

	out, err := runMirror(t, "--dir", root, "import-ollama", "--ollama-dir", ollama)
	if err != nil || !strings.Contains(out, "  qwen2.5:7b\n") {
		t.Fatalf("listing: %v\n%s", err, out)
	}

	out, err = runMirror(t, "--dir", root, "import-ollama", "qwen2.5:7b", "--ollama-dir", ollama,
		"--base", "http://models.internal:8081")
	if err != nil {
		t.Fatalf("%v\n%s", err, out)
	}
	file := filepath.Join(root, "qwen2.5", "7b-"+sum[:12], "qwen2.5-7b.gguf")
	if b, err := os.ReadFile(file); err != nil || !bytes.Equal(b, weights) {
		t.Fatalf("the GGUF is not in the mirror at %s: %v", file, err)
	}
	for _, want := range []string{
		"sha256 " + sum,
		"source: http://models.internal:8081/qwen2.5/7b-" + sum[:12] + "/qwen2.5-7b.gguf",
		"engine: llamacpp",
		`args: ["--ctx-size", "16384"]`,
	} {
		if !strings.Contains(out, want) {
			t.Errorf("the output does not say %q:\n%s", want, out)
		}
	}
}

func TestTheConsoleImportsFromOllamaIntoTheMirrorBesideIt(t *testing.T) {
	weights := []byte("GGUF\x03\x00\x00\x00llama weights")
	ollama, sum := ollamaWith(t, "llama3.2", "3b", weights, `{}`)
	root := t.TempDir()

	m, _ := consoleWith(t, cMirror)
	var tm tea.Model = press(m, "i")
	f := tm.(consoleModel).formD
	if f == nil || f.title != "Import from Ollama" {
		t.Fatalf("I opened %+v", f)
	}
	f.fields[0].text.SetValue("llama3.2:3b")
	f.fields[1].text.SetValue(ollama)
	f.fields[2].text.SetValue(root)
	for range len(f.fields) {
		tm = press(tm, "tab")
	}
	tm = pressRun(tm, "enter")

	if _, err := os.Stat(filepath.Join(root, "llama3.2", "3b-"+sum[:12], "llama3.2-3b.gguf")); err != nil {
		t.Fatalf("not imported: %v", err)
	}
	text := tm.(consoleModel).dialogText
	if !strings.Contains(text, "engine: llamacpp") || strings.Contains(text, "(100%)") {
		t.Errorf("the console said:\n%s", text)
	}
}

func TestImportCopiesALocalGGUFIntoTheMirror(t *testing.T) {
	src := filepath.Join(t.TempDir(), "q4_k_m.gguf")
	content := []byte("GGUF\x03\x00\x00\x00converted")
	if err := os.WriteFile(src, content, 0o644); err != nil {
		t.Fatal(err)
	}
	root := t.TempDir()
	out, err := runMirror(t, "--dir", root, "import", src, "--model", "qwen", "--version", "2026-10")
	if err != nil {
		t.Fatalf("%v\n%s", err, out)
	}
	if b, err := os.ReadFile(filepath.Join(root, "qwen", "2026-10", "q4_k_m.gguf")); err != nil || !bytes.Equal(b, content) {
		t.Fatalf("not imported: %v", err)
	}
	if !strings.Contains(out, "source: http://<this-host>:8081/qwen/2026-10/q4_k_m.gguf") {
		t.Errorf("no source given:\n%s", out)
	}
	// The layout is the point: a version is never replaced.
	if _, err := runMirror(t, "--dir", root, "import", src, "--model", "qwen", "--version", "2026-10"); err == nil {
		t.Error("an existing version was replaced")
	}
	if _, err := runMirror(t, "--dir", root, "import", src); err == nil {
		t.Error("imported without --model")
	}
}