god-eye/internal/ai/ensure_test.go

package ai

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func TestAlreadyInstalled(t *testing.T) {
	installed := map[string]bool{
		"qwen3:1.7b":            true,
		"qwen2.5-coder:14b":     true,
		"custom-model:latest":   true,
	}
	cases := []struct {
		model string
		want  bool
	}{
		{"qwen3:1.7b", true},
		{"qwen2.5-coder:14b", true},
		{"custom-model", true}, // via :latest fallback
		{"llama3:8b", false},
		{"qwen3", false}, // bare name: only matches when ":latest" is installed (it isn't)
	}
	for _, c := range cases {
		if got := alreadyInstalled(installed, c.model); got != c.want {
			t.Errorf("alreadyInstalled(%q) = %v, want %v", c.model, got, c.want)
		}
	}
}

func TestDedup(t *testing.T) {
	got := dedup([]string{"a", "b", "a", "", "c", " b "})
	want := []string{"a", "b", "c"}
	if len(got) != len(want) {
		t.Fatalf("got %v, want %v", got, want)
	}
	for i := range got {
		if got[i] != want[i] {
			t.Errorf("index %d: got %q want %q", i, got[i], want[i])
		}
	}
}

func TestHumanBytes(t *testing.T) {
	cases := []struct {
		in   int64
		want string
	}{
		{0, "0B"},
		{512, "512B"},
		{1024, "1.0KB"},
		{1024 * 1024, "1.0MB"},
		{1024 * 1024 * 1024, "1.0GB"},
		{int64(2.5 * 1024 * 1024 * 1024), "2.5GB"},
	}
	for _, c := range cases {
		if got := humanBytes(c.in); got != c.want {
			t.Errorf("humanBytes(%d) = %q, want %q", c.in, got, c.want)
		}
	}
}

func TestInstalled_ParsesTagsResponse(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/api/tags" {
			http.NotFound(w, r)
			return
		}
		_ = json.NewEncoder(w).Encode(map[string]interface{}{
			"models": []map[string]string{
				{"name": "qwen3:1.7b"},
				{"name": "qwen2.5-coder:14b"},
			},
		})
	}))
	defer srv.Close()

	e := NewModelEnsurer(srv.URL)
	got, err := e.Installed(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if !got["qwen3:1.7b"] || !got["qwen2.5-coder:14b"] {
		t.Errorf("missing expected models: %v", got)
	}
}

func TestInstalled_Non200ReturnsError(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "nope", http.StatusInternalServerError)
	}))
	defer srv.Close()

	e := NewModelEnsurer(srv.URL)
	if _, err := e.Installed(context.Background()); err == nil {
		t.Error("expected error on non-200")
	}
}

func TestReachable(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"models":[]}`))
	}))
	defer srv.Close()

	e := NewModelEnsurer(srv.URL)
	if err := e.Reachable(context.Background()); err != nil {
		t.Errorf("expected reachable, got %v", err)
	}
}

func TestReachable_Unreachable(t *testing.T) {
	e := NewModelEnsurer("http://127.0.0.1:1") // nothing listens here
	if err := e.Reachable(context.Background()); err == nil {
		t.Error("expected unreachable error")
	}
}

func TestPull_StreamsProgress(t *testing.T) {
	// Fake Ollama that emits a few NDJSON status events.
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/api/pull" {
			http.NotFound(w, r)
			return
		}
		w.Header().Set("Content-Type", "application/x-ndjson")
		events := []string{
			`{"status":"pulling manifest"}`,
			`{"status":"downloading","digest":"sha256:abc","total":1048576,"completed":524288}`,
			`{"status":"downloading","digest":"sha256:abc","total":1048576,"completed":1048576}`,
			`{"status":"verifying sha256 digest"}`,
			`{"status":"writing manifest"}`,
			`{"status":"success"}`,
		}
		for _, e := range events {
			w.Write([]byte(e + "\n"))
			if flusher, ok := w.(http.Flusher); ok {
				flusher.Flush()
			}
		}
	}))
	defer srv.Close()

	buf := &bytes.Buffer{}
	e := NewModelEnsurer(srv.URL)
	e.Verbose = true
	e.Writer = buf

	if err := e.Pull(context.Background(), "fake:1b"); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	out := buf.String()
	if !strings.Contains(out, "pulling manifest") {
		t.Errorf("missing 'pulling manifest' in output: %q", out)
	}
	if !strings.Contains(out, "success") {
		t.Errorf("missing 'success' in output: %q", out)
	}
}

func TestPull_ErrorBubblesUp(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"error":"model not found"}` + "\n"))
	}))
	defer srv.Close()

	e := NewModelEnsurer(srv.URL)
	err := e.Pull(context.Background(), "nonexistent")
	if err == nil {
		t.Fatal("expected error")
	}
	if !strings.Contains(err.Error(), "model not found") {
		t.Errorf("unexpected error: %v", err)
	}
}

func TestEnsureAll_SkipsInstalled_PullsMissing(t *testing.T) {
	pullCalls := map[string]int{}
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/api/tags":
			_ = json.NewEncoder(w).Encode(map[string]interface{}{
				"models": []map[string]string{{"name": "already-here:1b"}},
			})
		case "/api/pull":
			var body struct {
				Name string `json:"name"`
			}
			_ = json.NewDecoder(r.Body).Decode(&body)
			pullCalls[body.Name]++
			w.Write([]byte(`{"status":"success"}` + "\n"))
		default:
			http.NotFound(w, r)
		}
	}))
	defer srv.Close()

	e := NewModelEnsurer(srv.URL)
	if err := e.EnsureAll(context.Background(), []string{"already-here:1b", "missing-a:7b", "missing-b:14b"}); err != nil {
		t.Fatal(err)
	}
	if pullCalls["already-here:1b"] > 0 {
		t.Errorf("should not have pulled already-here")
	}
	if pullCalls["missing-a:7b"] != 1 || pullCalls["missing-b:14b"] != 1 {
		t.Errorf("missing models not pulled correctly: %v", pullCalls)
	}
}