golangLAKEHOUSE/internal/validator/iterate_test.go

package validator

import (
	"context"
	"errors"
	"testing"
)

func TestExtractJSON_FromFencedBlock(t *testing.T) {
	raw := "Here's my answer:\n```json\n{\"fills\": [{\"candidate_id\": \"W-1\"}]}\n```\nDone."
	v := ExtractJSON(raw)
	if v == nil {
		t.Fatal("expected match in fenced block")
	}
	if _, ok := v["fills"]; !ok {
		t.Errorf("missing fills key: %+v", v)
	}
}

func TestExtractJSON_FromBareBraces(t *testing.T) {
	raw := "Here you go: {\"fills\": [{\"candidate_id\": \"W-2\"}]}"
	v := ExtractJSON(raw)
	if v == nil {
		t.Fatal("expected match in bare braces")
	}
}

func TestExtractJSON_ReturnsNilOnNoObject(t *testing.T) {
	if v := ExtractJSON("just prose, no json"); v != nil {
		t.Errorf("expected nil, got %+v", v)
	}
}

func TestExtractJSON_PicksFirstBalancedObject(t *testing.T) {
	v := ExtractJSON(`{"a":1} then {"b":2}`)
	if v == nil {
		t.Fatal("expected match")
	}
	if v["a"].(float64) != 1 {
		t.Errorf("expected first object, got %+v", v)
	}
}

func TestExtractJSON_NestedBalancedObjects(t *testing.T) {
	v := ExtractJSON(`prefix {"outer": {"inner": [1,2,3]}, "x": "y"} suffix`)
	if v == nil {
		t.Fatal("expected match on balanced nested object")
	}
	if outer, ok := v["outer"].(map[string]any); !ok || outer["inner"] == nil {
		t.Errorf("nested structure lost: %+v", v)
	}
}

func TestExtractJSON_TopLevelArrayReturnsFirstInnerObject(t *testing.T) {
	// Both Rust and Go runtimes accept the first balanced {...} as a
	// successful match — for `[{"a":1},{"b":2}]` that's the first
	// inner object. Documenting this so the contract stays consistent
	// across runtimes.
	v := ExtractJSON(`[{"a":1},{"b":2}]`)
	if v == nil {
		t.Fatal("expected first inner object to be returned")
	}
	if v["a"].(float64) != 1 {
		t.Errorf("expected first object {a:1}, got %+v", v)
	}
}

// ─── Iterate orchestrator tests with scripted ChatCaller ────────────

func scriptedChat(responses ...string) (ChatCaller, *int) {
	idx := 0
	return func(_ context.Context, _, _ string, _, _ string, _ *float64, _ int) (string, error) {
		if idx >= len(responses) {
			return "", errors.New("scripted chat exhausted")
		}
		r := responses[idx]
		idx++
		return r, nil
	}, &idx
}

func TestIterate_AcceptsFirstValidArtifact(t *testing.T) {
	chat, calls := scriptedChat(`{"endorsed_names":["W-1"]}`)
	validate := func(_ string, _ map[string]any) (Report, error) {
		return Report{ElapsedMs: 1}, nil
	}
	resp, fail, err := Iterate(context.Background(),
		IterateRequest{Kind: "playbook", Prompt: "produce X", Provider: "ollama", Model: "qwen3.5:latest"},
		IterateConfig{}, chat, validate)
	if err != nil || fail != nil {
		t.Fatalf("expected success, got err=%v fail=%+v", err, fail)
	}
	if resp.Iterations != 1 {
		t.Errorf("iterations = %d, want 1", resp.Iterations)
	}
	if len(resp.History) != 1 || resp.History[0].Status.Kind != "accepted" {
		t.Errorf("history: %+v", resp.History)
	}
	if *calls != 1 {
		t.Errorf("expected 1 chat call, got %d", *calls)
	}
}

func TestIterate_RetriesOnNoJsonThenSucceeds(t *testing.T) {
	chat, _ := scriptedChat(
		"sorry I cannot do that",
		`{"endorsed_names":["W-1"]}`,
	)
	validate := func(_ string, _ map[string]any) (Report, error) {
		return Report{}, nil
	}
	resp, _, err := Iterate(context.Background(),
		IterateRequest{Kind: "playbook", Prompt: "produce X", Provider: "ollama", Model: "x"},
		IterateConfig{}, chat, validate)
	if err != nil || resp == nil {
		t.Fatalf("expected success, err=%v", err)
	}
	if resp.Iterations != 2 {
		t.Errorf("iterations = %d, want 2", resp.Iterations)
	}
	if resp.History[0].Status.Kind != "no_json" {
		t.Errorf("first history status: %+v", resp.History[0].Status)
	}
}

func TestIterate_RetriesOnValidationFailureThenSucceeds(t *testing.T) {
	chat, _ := scriptedChat(
		`{"bad":"shape"}`,
		`{"good":"shape"}`,
	)
	calls := 0
	validate := func(_ string, body map[string]any) (Report, error) {
		calls++
		if _, ok := body["good"]; ok {
			return Report{}, nil
		}
		return Report{}, &ValidationError{Kind: ErrSchema, Field: "x", Reason: "missing good"}
	}
	resp, _, err := Iterate(context.Background(),
		IterateRequest{Kind: "playbook", Prompt: "produce X", Provider: "ollama", Model: "x"},
		IterateConfig{}, chat, validate)
	if err != nil || resp == nil {
		t.Fatalf("expected success, err=%v", err)
	}
	if calls != 2 {
		t.Errorf("validate calls = %d, want 2", calls)
	}
	if resp.History[0].Status.Kind != "validation_failed" {
		t.Errorf("first history status: %+v", resp.History[0].Status)
	}
	if resp.History[0].Status.Error == "" {
		t.Errorf("validation_failed entry must carry error string")
	}
}

func TestIterate_MaxIterationsExhaustedReturnsFailure(t *testing.T) {
	chat, _ := scriptedChat(`{}`, `{}`, `{}`)
	validate := func(_ string, _ map[string]any) (Report, error) {
		return Report{}, &ValidationError{Kind: ErrCompleteness, Reason: "always wrong"}
	}
	resp, fail, err := Iterate(context.Background(),
		IterateRequest{Kind: "playbook", Prompt: "X", Provider: "ollama", Model: "x", MaxIterations: 3},
		IterateConfig{}, chat, validate)
	if err != nil {
		t.Fatalf("infrastructure error unexpected: %v", err)
	}
	if resp != nil {
		t.Fatalf("expected failure, got %+v", resp)
	}
	if fail.Iterations != 3 {
		t.Errorf("iterations = %d, want 3", fail.Iterations)
	}
	if len(fail.History) != 3 {
		t.Errorf("history length = %d, want 3", len(fail.History))
	}
}

func TestIterate_PropagatesChatInfraError(t *testing.T) {
	chat := func(_ context.Context, _, _ string, _, _ string, _ *float64, _ int) (string, error) {
		return "", errors.New("connection refused")
	}
	validate := func(_ string, _ map[string]any) (Report, error) { return Report{}, nil }
	_, _, err := Iterate(context.Background(),
		IterateRequest{Kind: "playbook", Prompt: "X", Provider: "ollama", Model: "x"},
		IterateConfig{}, chat, validate)
	if err == nil {
		t.Fatal("expected infrastructure error to surface")
	}
}