iop/apps/node/internal/adapters/ollama/ollama_test.go

package ollama

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

	"go.uber.org/zap"

	noderuntime "iop/apps/node/internal/runtime"
	"iop/packages/config"
)

type fakeSink struct {
	mu     sync.Mutex
	events []noderuntime.RuntimeEvent
}

func (s *fakeSink) Emit(_ context.Context, event noderuntime.RuntimeEvent) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.events = append(s.events, event)
	return nil
}

func (s *fakeSink) all() []noderuntime.RuntimeEvent {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]noderuntime.RuntimeEvent(nil), s.events...)
}

func TestOllamaExecuteStreamsChatDeltas(t *testing.T) {
	var gotModel string
	var gotPrompt string
	var gotNumCtx int
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/api/chat" {
			t.Fatalf("unexpected path %s", r.URL.Path)
		}
		var req ollamaChatRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			t.Fatalf("decode request: %v", err)
		}
		gotModel = req.Model
		if req.Options != nil {
			gotNumCtx = req.Options.NumCtx
		}
		if len(req.Messages) != 1 {
			t.Fatalf("expected one message, got %+v", req.Messages)
		}
		gotPrompt = req.Messages[0].Content
		w.Header().Set("Content-Type", "application/x-ndjson")
		_, _ = w.Write([]byte(`{"message":{"role":"assistant","content":"hello "},"done":false}` + "\n"))
		_, _ = w.Write([]byte(`{"message":{"role":"assistant","content":"world"},"done":false}` + "\n"))
		_, _ = w.Write([]byte(`{"done":true,"prompt_eval_count":3,"eval_count":2}` + "\n"))
	}))
	defer server.Close()

	adapter := New(config.OllamaConf{BaseURL: server.URL, ContextSize: 262144}, zap.NewNop())
	sink := &fakeSink{}
	err := adapter.Execute(context.Background(), noderuntime.ExecutionSpec{
		RunID:  "run-1",
		Target: "llama-test",
		Input:  map[string]any{"prompt": "say hello"},
	}, sink)
	if err != nil {
		t.Fatalf("Execute failed: %v", err)
	}
	if gotModel != "llama-test" {
		t.Fatalf("model: got %q", gotModel)
	}
	if gotPrompt != "say hello" {
		t.Fatalf("prompt: got %q", gotPrompt)
	}
	if gotNumCtx != 262144 {
		t.Fatalf("num_ctx: got %d", gotNumCtx)
	}

	events := sink.all()
	if len(events) != 4 {
		t.Fatalf("expected 4 events, got %+v", events)
	}
	if events[0].Type != noderuntime.EventTypeStart {
		t.Fatalf("expected start event, got %+v", events[0])
	}
	if events[1].Delta+events[2].Delta != "hello world" {
		t.Fatalf("unexpected delta text: %+v", events)
	}
	if events[3].Type != noderuntime.EventTypeComplete {
		t.Fatalf("expected complete event, got %+v", events[3])
	}
	if events[3].Usage == nil || events[3].Usage.InputTokens != 3 || events[3].Usage.OutputTokens != 2 {
		t.Fatalf("unexpected usage: %+v", events[3].Usage)
	}
}

func TestOllamaExecuteEmitsErrorForHTTPFailure(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		http.Error(w, "boom", http.StatusBadGateway)
	}))
	defer server.Close()

	adapter := New(config.OllamaConf{BaseURL: server.URL}, zap.NewNop())
	sink := &fakeSink{}
	err := adapter.Execute(context.Background(), noderuntime.ExecutionSpec{
		RunID:  "run-err",
		Target: "llama-test",
		Input:  map[string]any{"prompt": "hi"},
	}, sink)
	if err == nil {
		t.Fatal("expected error")
	}
	events := sink.all()
	if len(events) < 2 || events[1].Type != noderuntime.EventTypeError {
		t.Fatalf("expected error event, got %+v", events)
	}
	if !strings.Contains(events[1].Error, "boom") {
		t.Fatalf("expected HTTP body in error, got %q", events[1].Error)
	}
}

func TestOllamaCapabilitiesQueryTags(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/api/tags" {
			t.Fatalf("unexpected path %s", r.URL.Path)
		}
		_, _ = w.Write([]byte(`{"models":[{"name":"llama-a"},{"name":"llama-b"}]}`))
	}))
	defer server.Close()

	adapter := New(config.OllamaConf{BaseURL: server.URL}, zap.NewNop())
	caps, err := adapter.Capabilities(context.Background())
	if err != nil {
		t.Fatalf("Capabilities failed: %v", err)
	}
	got := strings.Join(caps.Targets, ",")
	if got != "llama-a,llama-b" {
		t.Fatalf("targets: got %q", got)
	}
}