foreman/internal/ollama/inventory_test.go

package ollama

import (
	"context"
	"fmt"
	"log/slog"
	"net/http"
	"testing"
	"time"
)

// mockClient implements Client for inventory testing.
type mockClient struct {
	tagsFn func(ctx context.Context) (*TagsResponse, error)
	psFn   func(ctx context.Context) (*PsResponse, error)
}

func (m *mockClient) Chat(ctx context.Context, req ChatRequest, stream bool) (*ChatResponse, <-chan ChatResponse, error) {
	return nil, nil, fmt.Errorf("not implemented")
}

func (m *mockClient) Embed(ctx context.Context, req EmbedRequest) (*EmbedResponse, error) {
	return nil, fmt.Errorf("not implemented")
}

func (m *mockClient) Tags(ctx context.Context) (*TagsResponse, error) {
	return m.tagsFn(ctx)
}

func (m *mockClient) Ps(ctx context.Context) (*PsResponse, error) {
	return m.psFn(ctx)
}

func (m *mockClient) RawChat(ctx context.Context, body []byte) (*http.Response, error) {
	return nil, fmt.Errorf("not implemented")
}

func (m *mockClient) RawEmbed(ctx context.Context, body []byte) (*http.Response, error) {
	return nil, fmt.Errorf("not implemented")
}

func TestInventory_RefreshPopulatesModels(t *testing.T) {
	client := &mockClient{
		tagsFn: func(ctx context.Context) (*TagsResponse, error) {
			return &TagsResponse{
				Models: []ModelInfo{
					{Name: "qwen3:30b"},
					{Name: "nomic-embed-text"},
				},
			}, nil
		},
		psFn: func(ctx context.Context) (*PsResponse, error) {
			return &PsResponse{
				Models: []RunningModel{
					{Name: "nomic-embed-text"},
				},
			}, nil
		},
	}

	inv := NewModelInventory(client, slog.Default())
	if err := inv.Refresh(context.Background()); err != nil {
		t.Fatalf("Refresh: %v", err)
	}

	models := inv.Models()
	if len(models) != 2 {
		t.Fatalf("got %d models, want 2", len(models))
	}

	if !inv.HasModel("qwen3:30b") {
		t.Error("HasModel(qwen3:30b) = false, want true")
	}
	if inv.HasModel("nonexistent") {
		t.Error("HasModel(nonexistent) = true, want false")
	}

	resident := inv.ResidentModels()
	if len(resident) != 1 {
		t.Fatalf("got %d resident models, want 1", len(resident))
	}

	if inv.Degraded() {
		t.Error("degraded should be false after successful refresh")
	}
	if inv.LastPoll().IsZero() {
		t.Error("lastPoll should be non-zero after refresh")
	}
}

func TestInventory_DegradedOnFailure(t *testing.T) {
	callCount := 0
	client := &mockClient{
		tagsFn: func(ctx context.Context) (*TagsResponse, error) {
			callCount++
			if callCount == 1 {
				return &TagsResponse{
					Models: []ModelInfo{{Name: "qwen3:30b"}},
				}, nil
			}
			return nil, fmt.Errorf("connection refused")
		},
		psFn: func(ctx context.Context) (*PsResponse, error) {
			return &PsResponse{}, nil
		},
	}

	inv := NewModelInventory(client, slog.Default())

	// First refresh succeeds.
	if err := inv.Refresh(context.Background()); err != nil {
		t.Fatalf("first Refresh: %v", err)
	}
	if inv.Degraded() {
		t.Error("should not be degraded after first successful poll")
	}

	// Second refresh fails — should retain models but mark degraded.
	if err := inv.Refresh(context.Background()); err == nil {
		t.Fatal("expected error on second refresh")
	}
	if !inv.Degraded() {
		t.Error("should be degraded after failed poll")
	}

	// Models should be retained.
	if !inv.HasModel("qwen3:30b") {
		t.Error("should retain models after failed poll")
	}
}

func TestInventory_RecoveryFromDegraded(t *testing.T) {
	failing := true
	client := &mockClient{
		tagsFn: func(ctx context.Context) (*TagsResponse, error) {
			if failing {
				return nil, fmt.Errorf("connection refused")
			}
			return &TagsResponse{
				Models: []ModelInfo{{Name: "qwen3:30b"}},
			}, nil
		},
		psFn: func(ctx context.Context) (*PsResponse, error) {
			return &PsResponse{}, nil
		},
	}

	inv := NewModelInventory(client, slog.Default())

	// First refresh fails.
	inv.Refresh(context.Background())
	if !inv.Degraded() {
		t.Error("should be degraded after failed poll")
	}

	// Target recovers.
	failing = false
	if err := inv.Refresh(context.Background()); err != nil {
		t.Fatalf("recovery Refresh: %v", err)
	}
	if inv.Degraded() {
		t.Error("should not be degraded after successful poll")
	}
}

func TestInventory_StartAndCancel(t *testing.T) {
	pollCount := 0
	client := &mockClient{
		tagsFn: func(ctx context.Context) (*TagsResponse, error) {
			pollCount++
			return &TagsResponse{}, nil
		},
		psFn: func(ctx context.Context) (*PsResponse, error) {
			return &PsResponse{}, nil
		},
	}

	inv := NewModelInventory(client, slog.Default())
	ctx, cancel := context.WithCancel(context.Background())

	done := make(chan struct{})
	go func() {
		inv.Start(ctx, 10*time.Millisecond)
		close(done)
	}()

	// Let it poll a few times.
	time.Sleep(50 * time.Millisecond)
	cancel()

	select {
	case <-done:
		// Clean exit.
	case <-time.After(2 * time.Second):
		t.Fatal("Start did not exit after context cancellation")
	}

	if pollCount < 2 {
		t.Errorf("poll count = %d, want >= 2 (initial + at least one tick)", pollCount)
	}
}