1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"mm/internal/config"
	"mm/internal/engine"
)

// storeServer is a llmman-style server holding `models`; /api/pull answers
// with `lines` (one JSON object each) and records what it was asked.
func storeServer(t *testing.T, models []string, lines []string, pulls *[]string) string {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/v1/models":
			data := []map[string]string{}
			for _, m := range models {
				data = append(data, map[string]string{"id": m})
			}
			_ = json.NewEncoder(w).Encode(map[string]any{"data": data})
		case "/api/pull":
			var body map[string]string
			_ = json.NewDecoder(r.Body).Decode(&body)
			*pulls = append(*pulls, body["model"])
			for _, l := range lines {
				_, _ = w.Write([]byte(l + "\n"))
			}
		default:
			http.NotFound(w, r)
		}
	}))
	t.Cleanup(srv.Close)
	return srv.URL + "/v1"
}

var pullOK = []string{
	`{"status":"pulling manifest"}`,
	`{"completed":0,"status":"pulling","total":2000000000}`,
	`{"completed":1000000000,"status":"pulling","total":2000000000}`,
	`{"completed":2000000000,"status":"pulling","total":2000000000}`,
	`{"status":"success"}`,
}

// llmmanEngine is an engine on the fake server, with autoPull set as asked.
func llmmanEngine(t *testing.T, baseURL string, autoPull bool, agents map[string]config.AgentRef) *engine.Engine {
	t.Helper()
	none := ""
	withConfig(t, func(c *config.Config) {
		c.Provider, c.Model, c.BaseURL, c.Fallback, c.AutoPull = "llmman", "smol", baseURL, &none, autoPull
		c.MMConfigs = agents
	})
	e, err := engine.New(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	return e
}

// Off, nothing is asked; on, the missing model is pulled with its progress,
// one line per new tenth off a terminal (0 % is the first line), and a
// closing line.
func TestAutoPullPullsTheMissingModel(t *testing.T) {
	var pulls []string
	url := storeServer(t, nil, pullOK, &pulls)
	var out bytes.Buffer
	if err := autoPull(context.Background(), llmmanEngine(t, url, false, nil), &out, false); err != nil || len(pulls) != 0 || out.Len() != 0 {
		t.Fatalf("off: err %v, pulls %v, out %q", err, pulls, out.String())
	}
	if err := autoPull(context.Background(), llmmanEngine(t, url, true, nil), &out, false); err != nil {
		t.Fatal(err)
	}
	got := out.String()
	for _, want := range []string{"⬇ pulling smol\n", "⬇ pulling smol — 1.0 GB / 2.0 GB (50 %)\n", "⬇ pulling smol — 2.0 GB / 2.0 GB (100 %)\n", "⬇ pulled smol in "} {
		if !strings.Contains(got, want) {
			t.Errorf("missing %q in:\n%s", want, got)
		}
	}
	if len(pulls) != 1 || pulls[0] != "smol" {
		t.Errorf("pulls = %v", pulls)
	}
}

// The agent's own model failing stops mm with the server's words; an
// mmConfigs agent's model failing is a warning; a model present is silent.
func TestAutoPullFailures(t *testing.T) {
	var pulls []string
	bad := storeServer(t, nil, []string{`{"error":"HF model info: HTTP 401"}`}, &pulls)
	err := autoPull(context.Background(), llmmanEngine(t, bad, true, nil), &bytes.Buffer{}, false)
	if err == nil || !strings.Contains(err.Error(), "autoPull: smol: the server says: HF model info: HTTP 401") {
		t.Errorf("own model: %v", err)
	}

	ok := storeServer(t, []string{"smol"}, nil, &pulls)
	child := filepath.Join(t.TempDir(), "agent.go.yaml")
	if err := os.WriteFile(child, []byte("provider: llmman\nmodel: big\nbaseUrl: "+bad+"\nfallback: \"\"\n"), 0o644); err != nil {
		t.Fatal(err)
	}
	var out bytes.Buffer
	pulls = nil
	e := llmmanEngine(t, ok, true, map[string]config.AgentRef{"go": {Path: child}})
	if err := autoPull(context.Background(), e, &out, false); err != nil {
		t.Fatalf("an agent's model stopped mm: %v", err)
	}
	if !strings.Contains(out.String(), "[warning: autoPull: big ("+child+"): the server says: HF model info: HTTP 401]") {
		t.Errorf("out:\n%s", out.String())
	}
	if len(pulls) != 1 || pulls[0] != "big" {
		t.Errorf("pulls = %v (the present smol must not be pulled)", pulls)
	}
}

// On a terminal the progress is one line rewritten in place, cleared at the
// end.
func TestPullReportOnATerminal(t *testing.T) {
	var out bytes.Buffer
	r := newPullReport(&out, true)
	r.step(engine.PullProgress{Model: "m", Status: "pulling", Completed: 5e8, Total: 1e9})
	r.done("m", true, nil)
	if got := out.String(); !strings.HasPrefix(got, "\r⬇ pulling m — 0.5 GB / 1.0 GB (50 %)\033[K\r\033[K⬇ pulled m in ") {
		t.Errorf("out %q", got)
	}
}

// Sizes in GB from 1 GB, MB below; durations in tenths under ten seconds.
func TestPullReportUnits(t *testing.T) {
	if got := size(52_000_000, 145_000_000) + " / " + size(145_000_000, 145_000_000); got != "52 MB / 145 MB" {
		t.Errorf("small: %q", got)
	}
	if got := size(2_100_000_000, 8_100_000_000); got != "2.1 GB" {
		t.Errorf("large: %q", got)
	}
	if got := elapsed(840 * time.Millisecond); got != "0.8s" {
		t.Errorf("short: %q", got)
	}
	if got := elapsed(123400 * time.Millisecond); got != "2m3s" {
		t.Errorf("long: %q", got)
	}
}